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

dotnet/src/ZB.MOM.NatsNet.Server/JetStream/JetStreamEngine.cs

@@ -3,6 +3,7 @@ namespace ZB.MOM.NatsNet.Server;
 internal sealed class JetStreamEngine(JetStream state)
 {
     private readonly JetStream _state = state;
+    private static readonly TimeSpan MinUsageUpdateWindow = TimeSpan.FromMilliseconds(250);
 
     internal void SetStarted()
     {
@@ -44,4 +45,105 @@ internal sealed class JetStreamEngine(JetStream state)
             _state.Lock.ExitReadLock();
         }
     }
+
+    internal Exception? DisableJetStream(JsAccount? account)
+    {
+        if (account?.Account is not Account a)
+            return new InvalidOperationException("jetstream not enabled for account");
+
+        _state.Lock.EnterWriteLock();
+        try
+        {
+            _state.Accounts.Remove(a.Name);
+        }
+        finally
+        {
+            _state.Lock.ExitWriteLock();
+        }
+
+        account.Delete();
+        return null;
+    }
+
+    internal bool WouldExceedLimits(StorageType storageType, int size)
+    {
+        var total = storageType == StorageType.MemoryStorage
+            ? Interlocked.Read(ref _state.MemUsed)
+            : Interlocked.Read(ref _state.StoreUsed);
+        var max = storageType == StorageType.MemoryStorage
+            ? _state.Config.MaxMemory
+            : _state.Config.MaxStore;
+        return total + size > max;
+    }
+
+    internal bool LimitsExceeded(StorageType storageType) => WouldExceedLimits(storageType, 0);
+
+    internal static string TierName(int replicas) => $"R{(replicas <= 0 ? 1 : replicas)}";
+
+    internal static bool IsSameTier(StreamConfig cfgA, StreamConfig cfgB) =>
+        cfgA.Replicas == cfgB.Replicas;
+
+    internal static Dictionary<string, JetStreamAccountLimits> DiffCheckedLimits(
+        Dictionary<string, JetStreamAccountLimits> a,
+        Dictionary<string, JetStreamAccountLimits> b)
+    {
+        var diff = new Dictionary<string, JetStreamAccountLimits>(StringComparer.Ordinal);
+
+        foreach (var (tier, oldLimit) in a)
+        {
+            b.TryGetValue(tier, out var newLimit);
+            newLimit ??= new JetStreamAccountLimits();
+            diff[tier] = new JetStreamAccountLimits
+            {
+                MaxMemory = newLimit.MaxMemory - oldLimit.MaxMemory,
+                MaxStore = newLimit.MaxStore - oldLimit.MaxStore,
+            };
+        }
+
+        foreach (var (tier, newLimit) in b)
+        {
+            if (a.ContainsKey(tier))
+                continue;
+
+            diff[tier] = new JetStreamAccountLimits
+            {
+                MaxMemory = newLimit.MaxMemory,
+                MaxStore = newLimit.MaxStore,
+            };
+        }
+
+        return diff;
+    }
+
+    internal static (ulong Mem, ulong Store) ReservedStorage(
+        Dictionary<string, object?> streamAssignments,
+        string tier)
+    {
+        ulong mem = 0;
+        ulong store = 0;
+
+        foreach (var assignment in streamAssignments.Values.OfType<StreamAssignmentView>())
+        {
+            var cfg = assignment.Config;
+            if (!string.IsNullOrEmpty(tier) && !string.Equals(tier, 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;
+        }
+
+        return (mem, store);
+    }
+
+    internal static bool ShouldSendUsageUpdate(DateTime lastUpdateUtc) =>
+        DateTime.UtcNow - lastUpdateUtc >= MinUsageUpdateWindow;
+}
+
+internal sealed class StreamAssignmentView
+{
+    public required StreamConfig Config { get; init; }
 }