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

// Copyright 2012-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0

using System.Net;
using System.Text.Json;
using System.Threading.Channels;
using ZB.MOM.NatsNet.Server.Internal;

namespace ZB.MOM.NatsNet.Server;

public sealed partial class NatsServer
{
    private const int RaftPeerIdLength = 8;
    private const string PeerStateFileName = "peerstate.json";

    internal bool AccountNrgAllowed { get; set; } = true;

    internal Exception? BootstrapRaftNode(RaftConfig? cfg, IReadOnlyList<string>? knownPeers, bool allPeersKnown)
    {
        if (cfg is null)
        {
            return new InvalidOperationException("raft: nil config");
        }

        knownPeers ??= [];
        foreach (var peer in knownPeers)
        {
            if (string.IsNullOrWhiteSpace(peer) || peer.Length != RaftPeerIdLength)
            {
                return new InvalidOperationException($"raft: illegal peer: {peer}");
            }
        }

        var expected = knownPeers.Count;
        if (!allPeersKnown)
        {
            if (expected < 2)
            {
                expected = 2;
            }

            var opts = GetOpts();
            var routeCount = opts.Routes.Count;
            var gatewayPeerCount = 0;
            var clusterName = ClusterName();
            foreach (var gateway in opts.Gateway.Gateways)
            {
                if (string.Equals(gateway.Name, clusterName, StringComparison.Ordinal))
                {
                    continue;
                }

                foreach (var url in gateway.Urls)
                {
                    var host = url.Host;
                    if (IPAddress.TryParse(host, out _))
                    {
                        gatewayPeerCount++;
                    }
                    else
                    {
                        try
                        {
                            var addrs = Dns.GetHostAddresses(host);
                            gatewayPeerCount += addrs.Length > 0 ? addrs.Length : 1;
                        }
                        catch
                        {
                            gatewayPeerCount++;
                        }
                    }
                }
            }

            var inferred = routeCount + gatewayPeerCount;
            if (expected < inferred)
            {
                expected = inferred;
            }
        }

        if (string.IsNullOrWhiteSpace(cfg.Store))
        {
            return new InvalidOperationException("raft: storage directory is not set");
        }

        try
        {
            Directory.CreateDirectory(cfg.Store);

            var tmpPath = Path.Combine(cfg.Store, $"_test_{Guid.NewGuid():N}");
            using (File.Create(tmpPath)) { }
            File.Delete(tmpPath);

            var peerState = new RaftPeerState
            {
                KnownPeers = [.. knownPeers],
                ClusterSize = expected,
                DomainExt = 0,
            };

            var peerStatePath = Path.Combine(cfg.Store, PeerStateFileName);
            var json = JsonSerializer.Serialize(peerState);
            File.WriteAllText(peerStatePath, json);
        }
        catch (Exception ex)
        {
            return ex;
        }

        return null;
    }

    internal (Raft? Node, Exception? Error) InitRaftNode(string accName, RaftConfig? cfg, IReadOnlyDictionary<string, string>? labels = null)
    {
        if (cfg is null)
        {
            return (null, new InvalidOperationException("raft: nil config"));
        }

        _mu.EnterReadLock();
        try
        {
            if (_sys == null)
            {
                return (null, ServerErrors.ErrNoSysAccount);
            }
        }
        finally
        {
            _mu.ExitReadLock();
        }

        var node = new Raft
        {
            Created_ = DateTime.UtcNow,
            GroupName = cfg.Name,
            StoreDir = cfg.Store,
            Wal = cfg.Log,
            Track = cfg.Track,
            Observer_ = cfg.Observer,
            Initializing = !cfg.Recovering,
            ScaleUp_ = cfg.ScaleUp,
            AccName = accName,
            Server_ = this,
            Id = (ServerName() ?? string.Empty).PadRight(RaftPeerIdLength, '0')[..RaftPeerIdLength],
            Qn = 1,
            Csz = 1,
            StateValue = (int)RaftState.Follower,
            LeadC = Channel.CreateUnbounded<bool>(),
            Quit = Channel.CreateUnbounded<bool>(),
            ApplyQ_ = new IpQueue<CommittedEntry>($"{cfg.Name}-committed"),
            PropQ = new IpQueue<ProposedEntry>($"{cfg.Name}-propose"),
            EntryQ = new IpQueue<AppendEntry>($"{cfg.Name}-append"),
            RespQ = new IpQueue<AppendEntryResponse>($"{cfg.Name}-append-response"),
            Reqs = new IpQueue<VoteRequest>($"{cfg.Name}-vote-req"),
            Votes_ = new IpQueue<VoteResponse>($"{cfg.Name}-vote-resp"),
        };

        RegisterRaftNode(node.GroupName, node);
        _ = labels;
        return (node, null);
    }

    internal (IRaftNode? Node, Exception? Error) StartRaftNode(string accName, RaftConfig? cfg, IReadOnlyDictionary<string, string>? labels = null)
    {
        var (node, error) = InitRaftNode(accName, cfg, labels);
        if (error is not null)
        {
            return (null, error);
        }

        return (node, null);
    }

    internal string ServerNameForNode(string node)
    {
        if (_nodeToInfo.TryGetValue(node, out var value) && value is NodeInfo info)
        {
            return info.Name;
        }

        return string.Empty;
    }

    internal string ClusterNameForNode(string node)
    {
        if (_nodeToInfo.TryGetValue(node, out var value) && value is NodeInfo info)
        {
            return info.Cluster;
        }

        return string.Empty;
    }

    internal void RegisterRaftNode(string group, IRaftNode node)
    {
        _raftNodes[group] = node;
    }

    internal void UnregisterRaftNode(string group)
    {
        _raftNodes.TryRemove(group, out _);
    }

    internal int NumRaftNodes() => _raftNodes.Count;

    internal IRaftNode? LookupRaftNode(string group)
    {
        if (_raftNodes.TryGetValue(group, out var value) && value is IRaftNode node)
        {
            return node;
        }

        return null;
    }

    internal void ReloadDebugRaftNodes(bool debug)
    {
        foreach (var value in _raftNodes.Values)
        {
            if (value is Raft raft)
            {
                raft.DebugEnabled = debug;
            }
        }
    }

    internal NodeInfo? GetNodeInfo(string nodeId)
    {
        if (_nodeToInfo.TryGetValue(nodeId, out var value) && value is NodeInfo info)
        {
            return info;
        }

        return null;
    }
}

internal sealed class RaftPeerState
{
    public List<string> KnownPeers { get; set; } = [];
    public int ClusterSize { get; set; }
    public ushort DomainExt { get; set; }
}