harmonia/Harmonia.Core/Caching/Cache.cs

using System.Collections.Concurrent;

namespace Harmonia.Core.Caching;

public abstract class Cache<TKey, TValue> : ICache<TKey, TValue> where TKey : notnull
{
    private readonly ConcurrentDictionary<object, SemaphoreSlim> _locks;
    private readonly SemaphoreSlim _throttler;

    protected abstract int MaxConcurrentRequests { get; }
    protected abstract TValue? TryGetValue(object key);
    protected abstract ValueTask<TValue?> FetchAsync(TKey key, CancellationToken cancellationToken);
    protected abstract void AddToCache(object key, TValue value);

    public Cache()
    {
        _locks = new();
        _throttler = new SemaphoreSlim(MaxConcurrentRequests, MaxConcurrentRequests);
    }

    public async Task<TValue?> GetAsync(TKey key, CancellationToken cancellationToken)
    {
        object? actualKey = GetKey(key);

        if (actualKey == null)
            return default;

        TValue? value = TryGetValue(actualKey);

        if (value != null)
            return value;

        SemaphoreSlim lockObject = _locks.GetOrAdd(actualKey, (key) => new SemaphoreSlim(1, 1));

        bool throttlerAcquired = false;
        bool lockAcquired = false;

        try
        {
            await _throttler.WaitAsync(cancellationToken);
            throttlerAcquired = true;

            if (cancellationToken.IsCancellationRequested)
                return default;

            await lockObject.WaitAsync(cancellationToken);
            lockAcquired = true;

            if (cancellationToken.IsCancellationRequested)
                return default;

            return TryGetValue(actualKey) ?? await FetchAsync2(key, cancellationToken);
        }
        catch (OperationCanceledException)
        {
            return default;
        }
        finally
        {
            if (lockAcquired)
                lockObject.Release();

            _locks.TryRemove(lockObject, out _);

            if (throttlerAcquired)
                _throttler.Release();
        }
    }

    private async Task<TValue?> FetchAsync2(TKey key, CancellationToken cancellationToken)
    {
        object? actualKey = GetKey(key);

        if (actualKey == null)
            return default;

        TValue? value = await FetchAsync(key, cancellationToken);

        if (value == null)
            return default;

        AddToCache(actualKey, value);

        return value;
    }

    public async Task<TValue?> RefreshAsync(TKey key, CancellationToken cancellationToken)
    {
        object? actualKey = GetKey(key);

        if (actualKey == null)
            return default;

        TValue? value = await FetchAsync(key, cancellationToken);

        if (value == null)
            return default;

        AddToCache(actualKey, value);

        return value;
    }

    protected virtual object? GetKey(TKey key)
    {
        return key;
    }
}