natsnet/dotnet/src/ZB.MOM.NatsNet.Server/JetStream/NatsStream.SnapshotMonitor.cs

namespace ZB.MOM.NatsNet.Server;

internal sealed partial class NatsStream
{
    private readonly object _preAcksSync = new();
    private readonly Dictionary<ulong, HashSet<NatsConsumer>> _preAcks = new();
    private bool _inMonitor;
    private long _replicationOutMsgs;
    private long _replicationOutBytes;

    internal bool NoInterestWithSubject(ulong seq, string subject, NatsConsumer? observingConsumer) =>
        !CheckForInterestWithSubject(seq, subject, observingConsumer);

    internal bool CheckForInterest(ulong seq, NatsConsumer? observingConsumer)
    {
        var subject = string.Empty;
        if (PotentialFilteredConsumers() && Store != null)
        {
            var loaded = Store.LoadMsg(seq, new StoreMsg());
            if (loaded == null)
            {
                RegisterPreAck(observingConsumer, seq);
                return true;
            }

            subject = loaded.Subject;
        }

        return CheckForInterestWithSubject(seq, subject, observingConsumer);
    }

    internal bool CheckForInterestWithSubject(ulong seq, string subject, NatsConsumer? observingConsumer)
    {
        _ = subject;
        lock (_consumersSync)
        {
            foreach (var consumer in _consumerList)
            {
                if (ReferenceEquals(consumer, observingConsumer))
                    continue;

                if (!HasPreAck(consumer, seq))
                    return true;
            }
        }

        ClearAllPreAcks(seq);
        return false;
    }

    internal bool HasPreAck(NatsConsumer? consumer, ulong seq)
    {
        if (consumer == null)
            return false;

        lock (_preAcksSync)
        {
            return _preAcks.TryGetValue(seq, out var consumers) && consumers.Contains(consumer);
        }
    }

    internal bool HasAllPreAcks(ulong seq, string subject)
    {
        lock (_preAcksSync)
        {
            if (!_preAcks.TryGetValue(seq, out var consumers) || consumers.Count == 0)
                return false;
        }

        return NoInterestWithSubject(seq, subject, null);
    }

    internal void ClearAllPreAcks(ulong seq)
    {
        lock (_preAcksSync)
        {
            _preAcks.Remove(seq);
        }
    }

    internal void ClearAllPreAcksBelowFloor(ulong floor)
    {
        lock (_preAcksSync)
        {
            var keys = _preAcks.Keys.Where(k => k < floor).ToArray();
            foreach (var key in keys)
                _preAcks.Remove(key);
        }
    }

    internal void RegisterPreAckLock(NatsConsumer? consumer, ulong seq)
    {
        _mu.EnterWriteLock();
        try
        {
            RegisterPreAck(consumer, seq);
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal void RegisterPreAck(NatsConsumer? consumer, ulong seq)
    {
        if (consumer == null)
            return;

        lock (_preAcksSync)
        {
            if (!_preAcks.TryGetValue(seq, out var consumers))
            {
                consumers = [];
                _preAcks[seq] = consumers;
            }

            consumers.Add(consumer);
        }
    }

    internal void ClearPreAck(NatsConsumer? consumer, ulong seq)
    {
        if (consumer == null)
            return;

        lock (_preAcksSync)
        {
            if (!_preAcks.TryGetValue(seq, out var consumers))
                return;

            consumers.Remove(consumer);
            if (consumers.Count == 0)
                _preAcks.Remove(seq);
        }
    }

    internal bool AckMsg(NatsConsumer? consumer, ulong seq)
    {
        if (seq == 0 || Store == null)
            return false;

        if (Config.Retention == RetentionPolicy.LimitsPolicy)
            return false;

        var state = new StreamState();
        Store.FastState(state);
        if (seq > state.LastSeq)
        {
            RegisterPreAck(consumer, seq);
            return true;
        }

        ClearPreAck(consumer, seq);
        if (seq < state.FirstSeq)
            return false;

        if (!NoInterest(seq, null))
            return false;

        if (!IsClustered())
        {
            var (removed, _) = Store.RemoveMsg(seq);
            return removed;
        }

        return true;
    }

    internal (SnapshotResult? Result, Exception? Error) Snapshot(TimeSpan deadline, bool checkMsgs, bool includeConsumers)
    {
        if (Store == null)
            return (null, new InvalidOperationException("store not initialized"));

        return Store.Snapshot(deadline, includeConsumers, checkMsgs);
    }

    internal void CheckForOrphanMsgs()
    {
        if (Store == null)
            return;

        var state = new StreamState();
        Store.FastState(state);
        ClearAllPreAcksBelowFloor(state.FirstSeq);
    }

    internal void CheckConsumerReplication()
    {
        if (Config.Retention != RetentionPolicy.InterestPolicy)
            return;

        lock (_consumersSync)
        {
            foreach (var consumer in _consumerList)
            {
                if (consumer.Config.Replicas == 0)
                    continue;

                if (consumer.Config.Replicas != Config.Replicas)
                    throw new InvalidOperationException("consumer replicas must match stream replicas for interest retention");
            }
        }
    }

    internal bool CheckInMonitor()
    {
        _mu.EnterWriteLock();
        try
        {
            if (_inMonitor)
                return true;

            _inMonitor = true;
            return false;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal void ClearMonitorRunning()
    {
        _mu.EnterWriteLock();
        try
        {
            _inMonitor = false;
            DeleteBatchApplyState();
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal bool IsMonitorRunning()
    {
        _mu.EnterReadLock();
        try
        {
            return _inMonitor;
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }

    internal void TrackReplicationTraffic(IRaftNode node, int size, int replicas)
    {
        if (!node.IsSystemAccount() || replicas <= 1)
            return;

        var additionalMsgs = replicas - 1;
        var additionalBytes = size * (replicas - 1);
        Interlocked.Add(ref _replicationOutMsgs, additionalMsgs);
        Interlocked.Add(ref _replicationOutBytes, additionalBytes);
    }
}