batch37 task6 implement group E pre-ack snapshot and restore

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

@@ -0,0 +1,259 @@
+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);
+    }
+}