using System.Collections.Concurrent;

using System.Security.Cryptography;

using System.Text;

namespace NpgsqlRest;

public interface IRoutineCache

{

bool Get(RoutineEndpoint endpoint, string key, out object? result);

/// <summary>

/// Stores or updates a cached value.

/// <paramref name="overrideExpiration"/> (when non-null) takes precedence over <see cref="RoutineEndpoint.CacheExpiresIn"/>

/// and is used by <see cref="CacheWhenRule"/> "Then" overrides to apply a per-request TTL.

/// </summary>

void AddOrUpdate(RoutineEndpoint endpoint, string key, object? value, TimeSpan? overrideExpiration = null);

bool Remove(string key);

/// <summary>

/// Atomic get-or-compute. On a cache hit the cached value is returned and <paramref name="factory"/>

/// is NOT invoked. On a miss the factory runs the underlying work (e.g. opening the connection,

/// executing SQL and serializing the response) and its result is stored before returning.

/// <para>

/// Backends that override this provide stampede protection: a burst of concurrent calls with the

/// same <paramref name="key"/> coalesce into a single factory invocation; the remaining callers

/// await the in-flight result instead of each hitting the database. The default implementation

/// below performs a plain probe-then-compute with no coalescing, so pre-existing custom

/// implementations keep working unchanged (they simply gain no stampede protection).

/// </para>

/// </summary>

/// <param name="endpoint">The endpoint whose response is being cached (carries the default TTL).</param>

/// <param name="key">The cache key; identical keys coalesce into one factory invocation.</param>

/// <param name="factory">Runs the underlying work when the cache is cold. Should respect the token.</param>

/// <param name="overrideExpiration">When non-null, the per-request TTL to use instead of the endpoint default.</param>

/// <param name="cancellationToken">Cancels this caller's wait; see backend notes for coalescing semantics.</param>

async ValueTask<object?> GetOrCreateAsync(

RoutineEndpoint endpoint,

string key,

Func<CancellationToken, ValueTask<object?>> factory,

TimeSpan? overrideExpiration = null,

CancellationToken cancellationToken = default)

{

if (Get(endpoint, key, out var existing))

{

return existing;

}

var result = await factory(cancellationToken);

AddOrUpdate(endpoint, key, result, overrideExpiration);

return result;

}

}

public static class CacheKeyHasher

{

/// <summary>

/// Computes a SHA256 hash of the cache key and returns it as a hex string.

/// The result is always 64 characters (256 bits / 4 bits per hex char).

/// </summary>

public static string ComputeHash(string key)

{

var hash = SHA256.HashData(Encoding.UTF8.GetBytes(key));

return Convert.ToHexString(hash);

}

/// <summary>

/// Returns the effective cache key based on the hashing configuration.

/// If UseHashedCacheKeys is true and the key length exceeds HashKeyThreshold, returns a hashed key.

/// Otherwise, returns the original key.

/// </summary>

public static string GetEffectiveKey(string key, CacheOptions options)

{

if (options.UseHashedCacheKeys && key.Length > options.HashKeyThreshold)

{

return ComputeHash(key);

}

return key;

}

}

public class RoutineCache : IRoutineCache

{

private class CacheEntry

{

public object? Value { get; set; }

public DateTime? ExpirationTime { get; set; }

public bool IsExpired => ExpirationTime.HasValue && DateTime.UtcNow > ExpirationTime.Value;

}

private static readonly ConcurrentDictionary<string, CacheEntry> _cache = new(StringComparer.Ordinal);

// In-flight factory invocations, keyed by effective (post-hash) cache key. Used to coalesce a burst

// of identical concurrent requests into a single factory run (stampede protection). The value is a

// Lazy so that even if ConcurrentDictionary.GetOrAdd runs its value-factory more than once under

// contention, only the stored Lazy's task is ever started — guaranteeing a single execution.

private static readonly ConcurrentDictionary<string, Lazy<Task<object?>>> _inflight = new(StringComparer.Ordinal);

private static Timer? _cleanupTimer;

private static CacheOptions _options = new();

public static void Start(NpgsqlRestOptions options)

{

_options = options.CacheOptions;

_cleanupTimer = new Timer(

_ => CleanupExpiredEntriesInternal(),

null,

TimeSpan.FromSeconds(options.CacheOptions.MemoryCachePruneIntervalSeconds),

TimeSpan.FromSeconds(options.CacheOptions.MemoryCachePruneIntervalSeconds));

}

public static void Shutdown()

{

_cleanupTimer?.Dispose();

_cache.Clear();

}

private static void CleanupExpiredEntriesInternal()

{

var expiredKeys = _cache

.Where(kvp => kvp.Value.IsExpired)

.Select(kvp => kvp.Key)

.ToList();

foreach (var key in expiredKeys)

{

_cache.TryRemove(key, out _);

}

}

public bool Get(RoutineEndpoint endpoint, string key, out object? result)

{

var effectiveKey = CacheKeyHasher.GetEffectiveKey(key, _options);

if (_cache.TryGetValue(effectiveKey, out var entry))

{

if (entry.IsExpired)

{

_cache.TryRemove(effectiveKey, out _);

result = null;

return false;

}

result = entry.Value;

return true;

}

result = null;

return false;

}

public void AddOrUpdate(RoutineEndpoint endpoint, string key, object? value, TimeSpan? overrideExpiration = null)

{

var effectiveKey = CacheKeyHasher.GetEffectiveKey(key, _options);

var ttl = overrideExpiration ?? endpoint.CacheExpiresIn;

var entry = new CacheEntry

{

Value = value,

ExpirationTime = ttl.HasValue ? DateTime.UtcNow + ttl.Value : null

};

_cache[effectiveKey] = entry;

}

public bool Remove(string key)

{

var effectiveKey = CacheKeyHasher.GetEffectiveKey(key, _options);

return _cache.TryRemove(effectiveKey, out _);

}

public async ValueTask<object?> GetOrCreateAsync(

RoutineEndpoint endpoint,

string key,

Func<CancellationToken, ValueTask<object?>> factory,

TimeSpan? overrideExpiration = null,

CancellationToken cancellationToken = default)

{

var effectiveKey = CacheKeyHasher.GetEffectiveKey(key, _options);

while (true)

{

cancellationToken.ThrowIfCancellationRequested();

if (Get(endpoint, key, out var cached))

{

return cached;

}

// Coalesce: the first caller for this key creates the in-flight task and runs the factory

// (with its own live connection/command and its own token); concurrent callers await it.

var lazy = _inflight.GetOrAdd(

effectiveKey,

_ => new Lazy<Task<object?>>(() => RunFactoryAsync(endpoint, key, effectiveKey, factory, overrideExpiration, cancellationToken)));

try

{

return await lazy.Value.WaitAsync(cancellationToken);

}

catch (OperationCanceledException)

{

if (cancellationToken.IsCancellationRequested)

{

// This caller itself was cancelled — propagate.

throw;

}

// The lead caller (whose token drove the shared run) was cancelled before this waiter.

// Its in-flight entry has already been removed; loop and retry — the value may now be

// cached, or this caller becomes the new lead and runs with its own live resources.

}

}

}

private async Task<object?> RunFactoryAsync(

RoutineEndpoint endpoint,

string key,

string effectiveKey,

Func<CancellationToken, ValueTask<object?>> factory,

TimeSpan? overrideExpiration,

CancellationToken cancellationToken)

{

try

{

var result = await factory(cancellationToken).ConfigureAwait(false);

AddOrUpdate(endpoint, key, result, overrideExpiration);

return result;

}

finally

{

// Always free the slot so a later miss (or a failed/cancelled run) can start fresh.

_inflight.TryRemove(effectiveKey, out _);

}

}

}