natsdotnet/src/NATS.Server/NatsServer.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using NATS.Server.Protocol;
using NATS.Server.Subscriptions;

namespace NATS.Server;

public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
{
    private readonly NatsOptions _options;
    private readonly ConcurrentDictionary<ulong, NatsClient> _clients = new();
    private readonly SubList _subList = new();
    private readonly ServerInfo _serverInfo;
    private Socket? _listener;
    private ulong _nextClientId;

    public SubList SubList => _subList;

    public NatsServer(NatsOptions options)
    {
        _options = options;
        _serverInfo = new ServerInfo
        {
            ServerId = Guid.NewGuid().ToString("N")[..20].ToUpperInvariant(),
            ServerName = options.ServerName ?? $"nats-dotnet-{Environment.MachineName}",
            Version = NatsProtocol.Version,
            Host = options.Host,
            Port = options.Port,
            MaxPayload = options.MaxPayload,
        };
    }

    public async Task StartAsync(CancellationToken ct)
    {
        _listener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        _listener.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
        _listener.Bind(new IPEndPoint(
            _options.Host == "0.0.0.0" ? IPAddress.Any : IPAddress.Parse(_options.Host),
            _options.Port));
        _listener.Listen(128);

        try
        {
            while (!ct.IsCancellationRequested)
            {
                var socket = await _listener.AcceptAsync(ct);
                var clientId = Interlocked.Increment(ref _nextClientId);

                var client = new NatsClient(clientId, socket, _options, _serverInfo);
                client.Router = this;
                _clients[clientId] = client;

                _ = RunClientAsync(client, ct);
            }
        }
        catch (OperationCanceledException) { }
    }

    private async Task RunClientAsync(NatsClient client, CancellationToken ct)
    {
        try
        {
            await client.RunAsync(ct);
        }
        catch (Exception)
        {
            // Client disconnected or errored
        }
        finally
        {
            RemoveClient(client);
        }
    }

    public void ProcessMessage(string subject, string? replyTo, ReadOnlyMemory<byte> headers,
        ReadOnlyMemory<byte> payload, NatsClient sender)
    {
        var result = _subList.Match(subject);

        // Deliver to plain subscribers
        foreach (var sub in result.PlainSubs)
        {
            if (sub.Client == null || sub.Client == sender && !(sender.ClientOpts?.Echo ?? true))
                continue;

            DeliverMessage(sub, subject, replyTo, headers, payload);
        }

        // Deliver to one member of each queue group (round-robin)
        foreach (var queueGroup in result.QueueSubs)
        {
            if (queueGroup.Length == 0) continue;

            // Simple round-robin -- pick based on total delivered across group
            var idx = Math.Abs((int)Interlocked.Increment(ref sender.OutMsgs)) % queueGroup.Length;
            // Undo the OutMsgs increment -- it will be incremented properly in SendMessageAsync
            Interlocked.Decrement(ref sender.OutMsgs);

            for (int attempt = 0; attempt < queueGroup.Length; attempt++)
            {
                var sub = queueGroup[(idx + attempt) % queueGroup.Length];
                if (sub.Client != null && (sub.Client != sender || (sender.ClientOpts?.Echo ?? true)))
                {
                    DeliverMessage(sub, subject, replyTo, headers, payload);
                    break;
                }
            }
        }
    }

    private static void DeliverMessage(Subscription sub, string subject, string? replyTo,
        ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload)
    {
        var client = sub.Client;
        if (client == null) return;

        // Check auto-unsub
        var count = Interlocked.Increment(ref sub.MessageCount);
        if (sub.MaxMessages > 0 && count > sub.MaxMessages)
            return;

        // Fire and forget -- deliver asynchronously
        _ = client.SendMessageAsync(subject, sub.Sid, replyTo, headers, payload, CancellationToken.None);
    }

    public void RemoveClient(NatsClient client)
    {
        _clients.TryRemove(client.Id, out _);
        client.RemoveAllSubscriptions(_subList);
    }

    public void Dispose()
    {
        _listener?.Dispose();
        foreach (var client in _clients.Values)
            client.Dispose();
        _subList.Dispose();
    }
}