natsdotnet/src/NATS.Server/Auth/PermissionLruCache.cs

namespace NATS.Server.Auth;

/// <summary>
/// Fixed-capacity LRU cache for permission results, supporting both PUB and SUB entries
/// with generation-based invalidation.
/// Lock-protected (per-client, low contention).
/// Reference: Go client.go maxPermCacheSize=128.
/// </summary>
public sealed class PermissionLruCache
{
    private readonly int _capacity;
    private readonly Dictionary<string, LinkedListNode<(string Key, bool Value)>> _map;
    private readonly LinkedList<(string Key, bool Value)> _list = new();
    private readonly object _lock = new();

    // Generation tracking: _generation is the authoritative counter (bumped by Invalidate),
    // _cacheGeneration is what the cache was last synced to. A mismatch on access triggers a clear.
    private long _generation;
    private long _cacheGeneration;

    public PermissionLruCache(int capacity = 128)
    {
        _capacity = capacity;
        _map = new Dictionary<string, LinkedListNode<(string Key, bool Value)>>(capacity, StringComparer.Ordinal);
    }

    /// <summary>
    /// The current generation counter. Incremented on each <see cref="Invalidate"/> call.
    /// </summary>
    public long Generation => Interlocked.Read(ref _generation);

    /// <summary>
    /// Bumps the generation counter so that the next TryGet or TryGetSub call detects
    /// the staleness and clears all cached entries.
    /// </summary>
    public void Invalidate() => Interlocked.Increment(ref _generation);

    // Ensures the cache is cleared if Invalidate was called since the last access.
    // Must be called inside _lock.
    private void EnsureFresh()
    {
        var gen = Interlocked.Read(ref _generation);
        if (gen != _cacheGeneration)
        {
            _map.Clear();
            _list.Clear();
            _cacheGeneration = gen;
        }
    }

    // ── PUB API (backward-compatible) ────────────────────────────────────────

    /// <summary>Looks up a PUB permission for <paramref name="key"/>.</summary>
    public bool TryGet(string key, out bool value)
    {
        var internalKey = "P:" + key;
        lock (_lock)
        {
            EnsureFresh();
            if (_map.TryGetValue(internalKey, out var node))
            {
                value = node.Value.Value;
                _list.Remove(node);
                _list.AddFirst(node);
                return true;
            }

            value = default;
            return false;
        }
    }

    /// <summary>Stores a PUB permission for <paramref name="key"/>.</summary>
    public void Set(string key, bool value)
    {
        var internalKey = "P:" + key;
        lock (_lock)
        {
            EnsureFresh();
            SetInternal(internalKey, value);
        }
    }

    // ── SUB API ───────────────────────────────────────────────────────────────

    /// <summary>Looks up a SUB permission for <paramref name="subject"/>.</summary>
    public bool TryGetSub(string subject, out bool value)
    {
        var internalKey = "S:" + subject;
        lock (_lock)
        {
            EnsureFresh();
            if (_map.TryGetValue(internalKey, out var node))
            {
                value = node.Value.Value;
                _list.Remove(node);
                _list.AddFirst(node);
                return true;
            }

            value = default;
            return false;
        }
    }

    /// <summary>Stores a SUB permission for <paramref name="subject"/>.</summary>
    public void SetSub(string subject, bool allowed)
    {
        var internalKey = "S:" + subject;
        lock (_lock)
        {
            EnsureFresh();
            SetInternal(internalKey, allowed);
        }
    }

    // ── Shared ────────────────────────────────────────────────────────────────

    public int Count
    {
        get
        {
            lock (_lock)
            {
                return _map.Count;
            }
        }
    }

    // Must be called inside _lock.
    private void SetInternal(string internalKey, bool value)
    {
        if (_map.TryGetValue(internalKey, out var existing))
        {
            _list.Remove(existing);
            existing.Value = (internalKey, value);
            _list.AddFirst(existing);
            return;
        }

        if (_map.Count >= _capacity)
        {
            var last = _list.Last!;
            _map.Remove(last.Value.Key);
            _list.RemoveLast();
        }

        var node = new LinkedListNode<(string Key, bool Value)>((internalKey, value));
        _list.AddFirst(node);
        _map[internalKey] = node;
    }
}