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

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();
        }
    }
}