natsdotnet/src/NATS.Server/Subscriptions/RouteResultCache.cs

namespace NATS.Server.Subscriptions;

/// <summary>
/// Per-account LRU cache for subscription match results.
/// Avoids repeated SubList.Match() calls for the same subject.
/// Uses atomic generation ID for bulk invalidation when subscriptions change.
/// Go reference: server/client.go routeCache, maxResultCacheSize=8192.
/// </summary>
public sealed class RouteResultCache
{
    private readonly int _capacity;
    private readonly Dictionary<string, LinkedListNode<(string Key, SubListResult Value)>> _map;
    private readonly LinkedList<(string Key, SubListResult Value)> _list = new();
    private readonly Lock _lock = new();
    private long _generation;
    private long _cacheGeneration;

    // Stats
    private long _hits;
    private long _misses;

    public RouteResultCache(int capacity = 8192)
    {
        _capacity = capacity;
        _map = new Dictionary<string, LinkedListNode<(string Key, SubListResult Value)>>(capacity, StringComparer.Ordinal);
    }

    public long Generation => Interlocked.Read(ref _generation);
    public long Hits => Interlocked.Read(ref _hits);
    public long Misses => Interlocked.Read(ref _misses);
    public int Count { get { lock (_lock) return _map.Count; } }

    /// <summary>
    /// Attempts to retrieve a cached result for the given subject.
    /// On generation mismatch (subscriptions changed since last cache fill),
    /// clears all entries and returns false.
    /// Updates LRU order on a hit.
    /// </summary>
    public bool TryGet(string subject, out SubListResult? result)
    {
        lock (_lock)
        {
            var currentGeneration = Interlocked.Read(ref _generation);
            if (_cacheGeneration != currentGeneration)
            {
                ClearUnderLock();
                _cacheGeneration = currentGeneration;
                result = null;
                Interlocked.Increment(ref _misses);
                return false;
            }

            if (_map.TryGetValue(subject, out var node))
            {
                // Move to front (most recently used)
                _list.Remove(node);
                _list.AddFirst(node);
                result = node.Value.Value;
                Interlocked.Increment(ref _hits);
                return true;
            }
        }

        result = null;
        Interlocked.Increment(ref _misses);
        return false;
    }

    /// <summary>
    /// Adds or updates a subject's cached result.
    /// If at capacity, evicts the least recently used entry (tail of list).
    /// </summary>
    public void Set(string subject, SubListResult result)
    {
        lock (_lock)
        {
            if (_map.TryGetValue(subject, out var existing))
            {
                // Update value in place and move to front
                _list.Remove(existing);
                var updated = new LinkedListNode<(string Key, SubListResult Value)>((subject, result));
                _list.AddFirst(updated);
                _map[subject] = updated;
                return;
            }

            // Evict LRU tail if at capacity
            if (_map.Count >= _capacity && _list.Last is { } tail)
            {
                _map.Remove(tail.Value.Key);
                _list.RemoveLast();
            }

            var node = new LinkedListNode<(string Key, SubListResult Value)>((subject, result));
            _list.AddFirst(node);
            _map[subject] = node;
        }
    }

    /// <summary>
    /// Increments the generation counter so that the next TryGet detects a mismatch
    /// and clears stale entries. This is called whenever subscriptions change.
    /// </summary>
    public void Invalidate()
    {
        Interlocked.Increment(ref _generation);
    }

    /// <summary>
    /// Immediately clears all cached entries under the lock.
    /// </summary>
    public void Clear()
    {
        lock (_lock)
        {
            ClearUnderLock();
        }
    }

    private void ClearUnderLock()
    {
        _map.Clear();
        _list.Clear();
    }
}