natsdotnet/src/NATS.Server/Events/InternalEventSystem.cs

using System.Collections.Concurrent;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Threading.Channels;
using Microsoft.Extensions.Logging;
using NATS.Server.Auth;
using NATS.Server.Subscriptions;

namespace NATS.Server.Events;

/// <summary>
/// Internal publish message queued for the send loop.
/// </summary>
public sealed class PublishMessage
{
    public InternalClient? Client { get; init; }
    public required string Subject { get; init; }
    public string? Reply { get; init; }
    public byte[]? Headers { get; init; }
    public object? Body { get; init; }
    public bool Echo { get; init; }
    public bool IsLast { get; init; }
}

/// <summary>
/// Internal received message queued for the receive loop.
/// </summary>
public sealed class InternalSystemMessage
{
    public required Subscription? Sub { get; init; }
    public required INatsClient? Client { get; init; }
    public required Account? Account { get; init; }
    public required string Subject { get; init; }
    public required string? Reply { get; init; }
    public required ReadOnlyMemory<byte> Headers { get; init; }
    public required ReadOnlyMemory<byte> Message { get; init; }
    public required SystemMessageHandler Callback { get; init; }
}

/// <summary>
/// Manages the server's internal event system with Channel-based send/receive loops.
/// Maps to Go's internal struct in events.go:124-147 and the goroutines
/// internalSendLoop (events.go:495) and internalReceiveLoop (events.go:476).
/// </summary>
public sealed class InternalEventSystem : IAsyncDisposable
{
    private readonly ILogger _logger;
    private readonly Channel<PublishMessage> _sendQueue;
    private readonly Channel<InternalSystemMessage> _receiveQueue;
    private readonly Channel<InternalSystemMessage> _receiveQueuePings;
    private readonly CancellationTokenSource _cts = new();

    private Task? _sendLoop;
    private Task? _receiveLoop;
    private Task? _receiveLoopPings;
    private NatsServer? _server;

    private ulong _sequence;
    private int _subscriptionId;
    private readonly ConcurrentDictionary<string, SystemMessageHandler> _callbacks = new();

    public Account SystemAccount { get; }
    public InternalClient SystemClient { get; }
    public string ServerHash { get; }

    public InternalEventSystem(Account systemAccount, InternalClient systemClient, string serverName, ILogger logger)
    {
        _logger = logger;
        SystemAccount = systemAccount;
        SystemClient = systemClient;

        // Hash server name for inbox routing (matches Go's shash)
        ServerHash = Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(serverName)))[..8].ToLowerInvariant();

        _sendQueue = Channel.CreateUnbounded<PublishMessage>(new UnboundedChannelOptions { SingleReader = true });
        _receiveQueue = Channel.CreateUnbounded<InternalSystemMessage>(new UnboundedChannelOptions { SingleReader = true });
        _receiveQueuePings = Channel.CreateUnbounded<InternalSystemMessage>(new UnboundedChannelOptions { SingleReader = true });
    }

    public void Start(NatsServer server)
    {
        _server = server;
        var ct = _cts.Token;
        _sendLoop = Task.Run(() => InternalSendLoopAsync(ct), ct);
        _receiveLoop = Task.Run(() => InternalReceiveLoopAsync(_receiveQueue, ct), ct);
        _receiveLoopPings = Task.Run(() => InternalReceiveLoopAsync(_receiveQueuePings, ct), ct);

        // Periodic stats publish every 10 seconds
        _ = Task.Run(async () =>
        {
            using var timer = new PeriodicTimer(TimeSpan.FromSeconds(10));
            while (await timer.WaitForNextTickAsync(ct))
            {
                PublishServerStats();
            }
        }, ct);
    }

    /// <summary>
    /// Registers system request-reply monitoring services for this server.
    /// Maps to Go's initEventTracking in events.go.
    /// Sets up handlers for $SYS.REQ.SERVER.{id}.VARZ, HEALTHZ, SUBSZ, STATSZ, IDZ
    /// and wildcard $SYS.REQ.SERVER.PING.* subjects.
    /// </summary>
    public void InitEventTracking(NatsServer server)
    {
        _server = server;
        var serverId = server.ServerId;

        // Server-specific monitoring services
        RegisterService(serverId, "VARZ", server.HandleVarzRequest);
        RegisterService(serverId, "HEALTHZ", server.HandleHealthzRequest);
        RegisterService(serverId, "SUBSZ", server.HandleSubszRequest);
        RegisterService(serverId, "STATSZ", server.HandleStatszRequest);
        RegisterService(serverId, "IDZ", server.HandleIdzRequest);

        // Wildcard ping services (all servers respond)
        SysSubscribe(string.Format(EventSubjects.ServerPing, "VARZ"), WrapRequestHandler(server.HandleVarzRequest));
        SysSubscribe(string.Format(EventSubjects.ServerPing, "HEALTHZ"), WrapRequestHandler(server.HandleHealthzRequest));
        SysSubscribe(string.Format(EventSubjects.ServerPing, "IDZ"), WrapRequestHandler(server.HandleIdzRequest));
        SysSubscribe(string.Format(EventSubjects.ServerPing, "STATSZ"), WrapRequestHandler(server.HandleStatszRequest));
    }

    private void RegisterService(string serverId, string name, Action<string, string?> handler)
    {
        var subject = string.Format(EventSubjects.ServerReq, serverId, name);
        SysSubscribe(subject, WrapRequestHandler(handler));
    }

    private SystemMessageHandler WrapRequestHandler(Action<string, string?> handler)
    {
        return (sub, client, acc, subject, reply, hdr, msg) =>
        {
            handler(subject, reply);
        };
    }

    /// <summary>
    /// Publishes a $SYS.SERVER.{id}.STATSZ message with current server statistics.
    /// Maps to Go's sendStatsz in events.go.
    /// Can be called manually for testing or is invoked periodically by the stats timer.
    /// </summary>
    public void PublishServerStats()
    {
        if (_server == null) return;

        var subject = string.Format(EventSubjects.ServerStats, _server.ServerId);
        var process = System.Diagnostics.Process.GetCurrentProcess();

        var statsMsg = new ServerStatsMsg
        {
            Server = _server.BuildEventServerInfo(),
            Stats = new ServerStatsData
            {
                Start = _server.StartTime,
                Mem = process.WorkingSet64,
                Cores = Environment.ProcessorCount,
                Connections = _server.ClientCount,
                TotalConnections = Interlocked.Read(ref _server.Stats.TotalConnections),
                Subscriptions = SystemAccount.SubList.Count,
                Sent = new DataStats
                {
                    Msgs = Interlocked.Read(ref _server.Stats.OutMsgs),
                    Bytes = Interlocked.Read(ref _server.Stats.OutBytes),
                },
                Received = new DataStats
                {
                    Msgs = Interlocked.Read(ref _server.Stats.InMsgs),
                    Bytes = Interlocked.Read(ref _server.Stats.InBytes),
                },
                InMsgs = Interlocked.Read(ref _server.Stats.InMsgs),
                OutMsgs = Interlocked.Read(ref _server.Stats.OutMsgs),
                InBytes = Interlocked.Read(ref _server.Stats.InBytes),
                OutBytes = Interlocked.Read(ref _server.Stats.OutBytes),
                SlowConsumers = Interlocked.Read(ref _server.Stats.SlowConsumers),
            },
        };

        Enqueue(new PublishMessage { Subject = subject, Body = statsMsg });
    }

    /// <summary>
    /// Creates a system subscription in the system account's SubList.
    /// Maps to Go's sysSubscribe in events.go:2796.
    /// </summary>
    public Subscription SysSubscribe(string subject, SystemMessageHandler callback)
    {
        var sid = Interlocked.Increment(ref _subscriptionId).ToString();
        var sub = new Subscription
        {
            Subject = subject,
            Sid = sid,
            Client = SystemClient,
        };

        // Store callback keyed by SID so multiple subscriptions work
        _callbacks[sid] = callback;

        // Set a single routing callback on the system client that dispatches by SID
        SystemClient.MessageCallback = (subj, s, reply, hdr, msg) =>
        {
            if (_callbacks.TryGetValue(s, out var cb))
            {
                _receiveQueue.Writer.TryWrite(new InternalSystemMessage
                {
                    Sub = sub,
                    Client = SystemClient,
                    Account = SystemAccount,
                    Subject = subj,
                    Reply = reply,
                    Headers = hdr,
                    Message = msg,
                    Callback = cb,
                });
            }
        };

        SystemAccount.SubList.Insert(sub);
        return sub;
    }

    /// <summary>
    /// Returns the next monotonically increasing sequence number for event ordering.
    /// </summary>
    public ulong NextSequence() => Interlocked.Increment(ref _sequence);

    /// <summary>
    /// Enqueue an internal message for publishing through the send loop.
    /// </summary>
    public void Enqueue(PublishMessage message)
    {
        _sendQueue.Writer.TryWrite(message);
    }

    /// <summary>
    /// The send loop: serializes messages and delivers them via the server's routing.
    /// Maps to Go's internalSendLoop in events.go:495-668.
    /// </summary>
    private async Task InternalSendLoopAsync(CancellationToken ct)
    {
        try
        {
            await foreach (var pm in _sendQueue.Reader.ReadAllAsync(ct))
            {
                try
                {
                    var seq = Interlocked.Increment(ref _sequence);

                    // Serialize body to JSON
                    byte[] payload;
                    if (pm.Body is byte[] raw)
                    {
                        payload = raw;
                    }
                    else if (pm.Body != null)
                    {
                        // Try source-generated context first, fall back to reflection-based for unknown types
                        var bodyType = pm.Body.GetType();
                        var typeInfo = EventJsonContext.Default.GetTypeInfo(bodyType);
                        payload = typeInfo != null
                            ? JsonSerializer.SerializeToUtf8Bytes(pm.Body, typeInfo)
                            : JsonSerializer.SerializeToUtf8Bytes(pm.Body, bodyType);
                    }
                    else
                    {
                        payload = [];
                    }

                    // Deliver via the system account's SubList matching
                    var result = SystemAccount.SubList.Match(pm.Subject);

                    foreach (var sub in result.PlainSubs)
                    {
                        sub.Client?.SendMessage(pm.Subject, sub.Sid, pm.Reply,
                            pm.Headers ?? ReadOnlyMemory<byte>.Empty,
                            payload);
                    }

                    foreach (var queueGroup in result.QueueSubs)
                    {
                        if (queueGroup.Length == 0) continue;
                        var sub = queueGroup[0]; // Simple pick for internal
                        sub.Client?.SendMessage(pm.Subject, sub.Sid, pm.Reply,
                            pm.Headers ?? ReadOnlyMemory<byte>.Empty,
                            payload);
                    }

                    if (pm.IsLast)
                        break;
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "Error in internal send loop processing message on {Subject}", pm.Subject);
                }
            }
        }
        catch (OperationCanceledException)
        {
            // Normal shutdown
        }
    }

    /// <summary>
    /// The receive loop: dispatches callbacks for internally-received messages.
    /// Maps to Go's internalReceiveLoop in events.go:476-491.
    /// </summary>
    private async Task InternalReceiveLoopAsync(Channel<InternalSystemMessage> queue, CancellationToken ct)
    {
        try
        {
            await foreach (var msg in queue.Reader.ReadAllAsync(ct))
            {
                try
                {
                    msg.Callback(msg.Sub, msg.Client, msg.Account, msg.Subject, msg.Reply, msg.Headers, msg.Message);
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "Error in internal receive loop processing {Subject}", msg.Subject);
                }
            }
        }
        catch (OperationCanceledException)
        {
            // Normal shutdown
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _cts.CancelAsync();
        _sendQueue.Writer.TryComplete();
        _receiveQueue.Writer.TryComplete();
        _receiveQueuePings.Writer.TryComplete();

        if (_sendLoop != null) await _sendLoop.WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
        if (_receiveLoop != null) await _receiveLoop.WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
        if (_receiveLoopPings != null) await _receiveLoopPings.WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

        _cts.Dispose();
    }
}