natsnet/dotnet/src/ZB.MOM.NatsNet.Server/NatsServer.JetStreamClusterMeta.cs

using System.Text.Json;

namespace ZB.MOM.NatsNet.Server;

public sealed partial class NatsServer
{
    internal (int JetStreamEnabled, int Total) TrackedJetStreamServers()
    {
        if (!IsRunning() || _sys == null)
            return (-1, -1);

        var js = 0;
        var total = 0;
        foreach (var value in _nodeToInfo.Values)
        {
            if (value is not NodeInfo info)
                continue;

            if (info.Js)
                js++;
            total++;
        }

        return (js, total);
    }

    internal (JetStream? JetStream, JetStreamCluster? Cluster) GetJetStreamCluster()
    {
        if (IsShuttingDown())
            return (null, null);

        var js = GetJetStreamState();
        if (js == null)
            return (null, null);

        return (js, js.Cluster as JetStreamCluster);
    }

    public bool JetStreamIsClustered()
    {
        var js = GetJetStreamState();
        return js != null && Interlocked.CompareExchange(ref js.Clustered, 0, 0) == 1;
    }

    public bool JetStreamIsLeader()
    {
        var (_, cluster) = GetJetStreamCluster();
        return cluster?.IsLeader() == true;
    }

    public bool JetStreamIsCurrent()
    {
        var js = GetJetStreamState();
        if (js == null)
            return false;

        js.Lock.EnterReadLock();
        try
        {
            if (js.Cluster is not JetStreamCluster cluster)
                return true;

            return cluster.Meta?.Current() == true;
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    public Exception? JetStreamSnapshotMeta()
    {
        var js = GetJetStreamState();
        if (js == null)
            return ToException(JsApiErrors.NewJSNotEnabledError());

        js.Lock.EnterReadLock();
        JetStreamCluster? cluster;
        IRaftNode? meta;
        try
        {
            cluster = js.Cluster as JetStreamCluster;
            meta = cluster?.Meta;
            if (cluster == null || meta == null)
                return ToException(JsApiErrors.NewJSClusterNotActiveError());
            if (!cluster.IsLeader())
                return ToException(JsApiErrors.NewJSClusterNotLeaderError());
        }
        finally
        {
            js.Lock.ExitReadLock();
        }

        var snapshot = JsonSerializer.SerializeToUtf8Bytes(cluster!.Streams);
        meta!.InstallSnapshot(snapshot, force: false);
        return null;
    }

    public Exception? JetStreamStepdownStream(string account, string stream)
    {
        var (_, cluster) = GetJetStreamCluster();
        if (cluster == null)
            return ToException(JsApiErrors.NewJSClusterNotActiveError());

        var (acc, error) = LookupAccount(account);
        if (error != null)
            return error;
        if (acc == null)
            return new InvalidOperationException("account not found");

        var (_, streamError) = acc.LookupStream(stream);
        if (streamError != null)
            return streamError;

        if (cluster.Streams.TryGetValue(account, out var assignments) &&
            assignments.TryGetValue(stream, out var assignment))
        {
            assignment.Group?.Node?.StepDown();
        }

        return null;
    }

    public Exception? JetStreamStepdownConsumer(string account, string stream, string consumer)
    {
        var (_, cluster) = GetJetStreamCluster();
        if (cluster == null)
            return ToException(JsApiErrors.NewJSClusterNotActiveError());

        var (acc, error) = LookupAccount(account);
        if (error != null)
            return error;
        if (acc == null)
            return new InvalidOperationException("account not found");

        var (_, streamError) = acc.LookupStream(stream);
        if (streamError != null)
            return streamError;

        if (!cluster.Streams.TryGetValue(account, out var assignments) ||
            !assignments.TryGetValue(stream, out var assignment) ||
            assignment.Consumers == null ||
            !assignment.Consumers.TryGetValue(consumer, out var consumerAssignment))
        {
            return ToException(JsApiErrors.NewJSConsumerNotFoundError());
        }

        consumerAssignment.Group?.Node?.StepDown();
        return null;
    }

    public Exception? JetStreamSnapshotStream(string account, string stream)
    {
        var (_, cluster) = GetJetStreamCluster();
        if (cluster == null)
            return ToException(JsApiErrors.NewJSClusterNotActiveError());

        var (acc, error) = LookupAccount(account);
        if (error != null)
            return error;
        if (acc == null)
            return new InvalidOperationException("account not found");

        var (natsStream, streamError) = acc.LookupStream(stream);
        if (streamError != null)
            return streamError;
        if (natsStream == null)
            return null;

        if (cluster.Streams.TryGetValue(account, out var assignments) &&
            assignments.TryGetValue(stream, out var assignment))
        {
            var snapshot = JsonSerializer.SerializeToUtf8Bytes(natsStream.State());
            assignment.Group?.Node?.InstallSnapshot(snapshot, force: false);
        }

        return null;
    }

    public string[] JetStreamClusterPeers()
    {
        var js = GetJetStreamState();
        if (js == null)
            return [];

        js.Lock.EnterReadLock();
        try
        {
            if (js.Cluster is not JetStreamCluster cluster || cluster.Meta == null || !cluster.IsLeader())
                return [];

            var names = new List<string>();
            foreach (var peer in cluster.Meta.Peers())
            {
                if (!_nodeToInfo.TryGetValue(peer.Id, out var value) || value is not NodeInfo info)
                    continue;
                if (info.Offline || !info.Js || info.Stats == null)
                    continue;

                names.Add(info.Name);
            }

            return [.. names];
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    public bool JetStreamIsStreamLeader(string account, string stream)
    {
        var (js, cluster) = GetJetStreamCluster();
        if (js == null || cluster == null)
            return false;

        js.Lock.EnterReadLock();
        try
        {
            return cluster.IsStreamLeader(account, stream);
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    public bool JetStreamIsStreamCurrent(string account, string stream)
    {
        var js = GetJetStreamState();
        if (js == null)
            return false;

        js.Lock.EnterReadLock();
        try
        {
            return (js.Cluster as JetStreamCluster)?.IsStreamCurrent(account, stream) ?? false;
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    public bool JetStreamIsConsumerLeader(string account, string stream, string consumer)
    {
        var (js, cluster) = GetJetStreamCluster();
        if (js == null || cluster == null)
            return false;

        js.Lock.EnterReadLock();
        try
        {
            return cluster.IsConsumerLeader(account, stream, consumer);
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    internal Exception? EnableJetStreamClustering()
    {
        if (!IsRunning())
            return null;

        var js = GetJetStream();
        if (js == null)
            return ToException(JsApiErrors.NewJSNotEnabledForAccountError());

        if (js.IsClustered())
            return null;

        return js.SetupMetaGroup();
    }

    public bool JetStreamIsStreamAssigned(string account, string stream)
    {
        var (js, cluster) = GetJetStreamCluster();
        if (js == null || cluster == null)
            return false;

        var (acc, _) = LookupAccount(account);
        if (acc == null)
            return false;

        js.Lock.EnterReadLock();
        try
        {
            return cluster.IsStreamAssigned(acc, stream);
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    private static Exception ToException(JsApiError error) =>
        new(error.Description ?? "jetstream error");
}