feat(batch27): implement jetstream usage foundations and limit helpers

dotnet/src/ZB.MOM.NatsNet.Server/JetStream/JsAccount.Core.cs

@@ -0,0 +1,194 @@
+namespace ZB.MOM.NatsNet.Server;
+
+internal sealed partial class JsAccount
+{
+    internal (ulong Mem, ulong Store) ReservedStorage(string tier)
+    {
+        ulong mem = 0;
+        ulong store = 0;
+
+        Lock.EnterReadLock();
+        try
+        {
+            foreach (var stream in Streams.Values.OfType<NatsStream>())
+            {
+                var cfg = stream.Config;
+                if (!string.IsNullOrEmpty(tier) && !string.Equals(tier, JetStreamEngine.TierName(cfg.Replicas), StringComparison.Ordinal))
+                    continue;
+                if (cfg.MaxBytes <= 0)
+                    continue;
+
+                if (cfg.Storage == StorageType.FileStorage)
+                    store += (ulong)cfg.MaxBytes;
+                else if (cfg.Storage == StorageType.MemoryStorage)
+                    mem += (ulong)cfg.MaxBytes;
+            }
+        }
+        finally
+        {
+            Lock.ExitReadLock();
+        }
+
+        return (mem, store);
+    }
+
+    internal void RemoteUpdateUsage(byte[] message)
+    {
+        if (message.Length < 16)
+            return;
+
+        UsageLock.EnterWriteLock();
+        try
+        {
+            if (!Usage.TryGetValue(string.Empty, out var usage))
+            {
+                usage = new JsaStorage();
+                Usage[string.Empty] = usage;
+            }
+
+            usage.Total.Mem = BitConverter.ToInt64(message, 0);
+            usage.Total.Store = BitConverter.ToInt64(message, 8);
+        }
+        finally
+        {
+            UsageLock.ExitWriteLock();
+        }
+    }
+
+    internal void CheckAndSyncUsage(string tier, StorageType storageType)
+    {
+        if (Interlocked.CompareExchange(ref Sync, 1, 0) != 0)
+            return;
+
+        try
+        {
+            Lock.EnterReadLock();
+            try
+            {
+                long total = 0;
+                foreach (var stream in Streams.Values.OfType<NatsStream>())
+                {
+                    if (!string.Equals(JetStreamEngine.TierName(stream.Config.Replicas), tier, StringComparison.Ordinal))
+                        continue;
+                    if (stream.Config.Storage != storageType)
+                        continue;
+                    total += (long)stream.State().Bytes;
+                }
+
+                UsageLock.EnterWriteLock();
+                try
+                {
+                    Usage.TryGetValue(tier, out var usage);
+                    usage ??= new JsaStorage();
+                    Usage[tier] = usage;
+
+                    if (storageType == StorageType.MemoryStorage)
+                    {
+                        usage.Local.Mem = total;
+                        usage.Total.Mem = total;
+                    }
+                    else
+                    {
+                        usage.Local.Store = total;
+                        usage.Total.Store = total;
+                    }
+                }
+                finally
+                {
+                    UsageLock.ExitWriteLock();
+                }
+            }
+            finally
+            {
+                Lock.ExitReadLock();
+            }
+        }
+        finally
+        {
+            Interlocked.Exchange(ref Sync, 0);
+        }
+    }
+
+    internal void UpdateUsage(string tier, StorageType storageType, long delta)
+    {
+        UsageLock.EnterWriteLock();
+        try
+        {
+            Usage.TryGetValue(tier, out var usage);
+            usage ??= new JsaStorage();
+            Usage[tier] = usage;
+
+            if (storageType == StorageType.MemoryStorage)
+            {
+                usage.Local.Mem += delta;
+                usage.Total.Mem += delta;
+            }
+            else
+            {
+                usage.Local.Store += delta;
+                usage.Total.Store += delta;
+            }
+        }
+        finally
+        {
+            UsageLock.ExitWriteLock();
+        }
+    }
+
+    internal void SendClusterUsageUpdateTimer()
+    {
+        UsageLock.EnterWriteLock();
+        try
+        {
+            SendClusterUsageUpdate();
+        }
+        finally
+        {
+            UsageLock.ExitWriteLock();
+        }
+    }
+
+    internal void SendClusterUsageUpdate()
+    {
+        var now = DateTime.UtcNow;
+        if (!JetStreamEngine.ShouldSendUsageUpdate(LUpdate))
+            return;
+        LUpdate = now;
+
+        // Cluster bus publish is wired in later cluster sessions.
+        UsageApi = ApiTotal;
+        UsageErr = ApiErrors;
+    }
+
+    internal (JetStream? JetStream, bool Clustered) JetStreamAndClustered()
+    {
+        Lock.EnterReadLock();
+        try
+        {
+            var js = Js as JetStream;
+            return (js, js?.Cluster != null);
+        }
+        finally
+        {
+            Lock.ExitReadLock();
+        }
+    }
+
+    internal Account? Acc() => Account as Account;
+
+    internal void Delete()
+    {
+        Lock.EnterWriteLock();
+        try
+        {
+            Streams.Clear();
+            Inflight.Clear();
+            UpdatesSub = null;
+            UpdatesPub = string.Empty;
+        }
+        finally
+        {
+            Lock.ExitWriteLock();
+        }
+    }
+}