natsdotnet/tests/NATS.Server.Benchmark.Tests/Transport/WebSocketPubSubTests.cs

using System.Net.WebSockets;
using System.Text;
using NATS.Client.Core;
using NATS.Server.Benchmark.Tests.Harness;
using NATS.Server.Benchmark.Tests.Infrastructure;
using Xunit.Abstractions;

namespace NATS.Server.Benchmark.Tests.Transport;

[Collection("Benchmark-WebSocket")]
public class WebSocketPubSubTests(WebSocketServerFixture fixture, ITestOutputHelper output)
{
    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task WsPubSub1To1_128B()
    {
        const int payloadSize = 128;
        const int messageCount = 5_000;

        var dotnetResult = await RunWsPubSub("WebSocket PubSub 1:1 (128B)", "DotNet", fixture.DotNetWsPort, fixture.CreateDotNetNatsClient, payloadSize, messageCount);

        if (fixture.GoAvailable)
        {
            var goResult = await RunWsPubSub("WebSocket PubSub 1:1 (128B)", "Go", fixture.GoWsPort, fixture.CreateGoNatsClient, payloadSize, messageCount);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task WsPubNoSub_128B()
    {
        const int payloadSize = 128;
        const int messageCount = 10_000;

        var dotnetResult = await RunWsPubOnly("WebSocket Pub-Only (128B)", "DotNet", fixture.DotNetWsPort, payloadSize, messageCount);

        if (fixture.GoAvailable)
        {
            var goResult = await RunWsPubOnly("WebSocket Pub-Only (128B)", "Go", fixture.GoWsPort, payloadSize, messageCount);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    private static async Task<BenchmarkResult> RunWsPubSub(string name, string serverType, int wsPort, Func<NatsConnection> createNatsClient, int payloadSize, int messageCount)
    {
        var payload = new byte[payloadSize];
        var subject = $"bench.ws.pubsub.{Guid.NewGuid():N}";

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));

        using var ws = new ClientWebSocket();
        await ws.ConnectAsync(new Uri($"ws://127.0.0.1:{wsPort}"), cts.Token);

        var reader = new WsLineReader(ws);

        // Read INFO
        await reader.ReadLineAsync(cts.Token);

        // Send CONNECT + SUB + PING
        await WsSend(ws, "CONNECT {\"verbose\":false,\"protocol\":1}\r\n", cts.Token);
        await WsSend(ws, $"SUB {subject} 1\r\n", cts.Token);
        await WsSend(ws, "PING\r\n", cts.Token);
        await WaitForPong(reader, cts.Token);

        // NATS publisher
        await using var natsPub = createNatsClient();
        await natsPub.ConnectAsync();
        await natsPub.PingAsync(cts.Token);

        var sw = System.Diagnostics.Stopwatch.StartNew();

        for (var i = 0; i < messageCount; i++)
            await natsPub.PublishAsync(subject, payload, cancellationToken: cts.Token);
        await natsPub.PingAsync(cts.Token);

        // Read all MSG responses from WebSocket
        var received = 0;
        while (received < messageCount)
        {
            var line = await reader.ReadLineAsync(cts.Token);
            if (line.StartsWith("MSG ", StringComparison.Ordinal))
            {
                await reader.ReadLineAsync(cts.Token);
                received++;
            }
        }

        sw.Stop();

        await ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, cts.Token);

        return new BenchmarkResult
        {
            Name = name,
            ServerType = serverType,
            TotalMessages = messageCount,
            TotalBytes = (long)messageCount * payloadSize,
            Duration = sw.Elapsed,
        };
    }

    private static async Task<BenchmarkResult> RunWsPubOnly(string name, string serverType, int wsPort, int payloadSize, int messageCount)
    {
        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));

        using var ws = new ClientWebSocket();
        await ws.ConnectAsync(new Uri($"ws://127.0.0.1:{wsPort}"), cts.Token);

        var reader = new WsLineReader(ws);

        // Read INFO
        await reader.ReadLineAsync(cts.Token);

        // Send CONNECT
        await WsSend(ws, "CONNECT {\"verbose\":false,\"protocol\":1}\r\n", cts.Token);
        await WsSend(ws, "PING\r\n", cts.Token);
        await WaitForPong(reader, cts.Token);

        // Build a PUB command with raw binary payload
        var subject = $"bench.ws.pubonly.{Guid.NewGuid():N}";
        var pubLine = $"PUB {subject} {payloadSize}\r\n";
        var pubPayload = new byte[payloadSize];
        var pubCmd = Encoding.ASCII.GetBytes(pubLine)
            .Concat(pubPayload)
            .Concat(Encoding.ASCII.GetBytes("\r\n"))
            .ToArray();

        var sw = System.Diagnostics.Stopwatch.StartNew();
        for (var i = 0; i < messageCount; i++)
            await ws.SendAsync(pubCmd, WebSocketMessageType.Binary, true, cts.Token);

        // Flush with PING/PONG
        await WsSend(ws, "PING\r\n", cts.Token);
        await WaitForPong(reader, cts.Token);
        sw.Stop();

        await ws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, cts.Token);

        return new BenchmarkResult
        {
            Name = name,
            ServerType = serverType,
            TotalMessages = messageCount,
            TotalBytes = (long)messageCount * payloadSize,
            Duration = sw.Elapsed,
        };
    }

    /// <summary>
    /// Reads lines until PONG is received, skipping any INFO lines
    /// (Go server sends a second INFO after CONNECT with connect_info:true).
    /// </summary>
    private static async Task WaitForPong(WsLineReader reader, CancellationToken ct)
    {
        while (true)
        {
            var line = await reader.ReadLineAsync(ct);
            if (line == "PONG")
                return;
        }
    }

    private static async Task WsSend(ClientWebSocket ws, string data, CancellationToken ct)
    {
        var bytes = Encoding.ASCII.GetBytes(data);
        await ws.SendAsync(bytes, WebSocketMessageType.Binary, true, ct);
    }

    /// <summary>
    /// Buffers incoming WebSocket frames and returns one NATS protocol line at a time.
    /// </summary>
    private sealed class WsLineReader(ClientWebSocket ws)
    {
        private readonly byte[] _recvBuffer = new byte[65536];
        private readonly StringBuilder _pending = new();

        public async Task<string> ReadLineAsync(CancellationToken ct)
        {
            while (true)
            {
                var full = _pending.ToString();
                var crlfIdx = full.IndexOf("\r\n", StringComparison.Ordinal);
                if (crlfIdx >= 0)
                {
                    var line = full[..crlfIdx];
                    _pending.Clear();
                    _pending.Append(full[(crlfIdx + 2)..]);
                    return line;
                }

                var result = await ws.ReceiveAsync(_recvBuffer, ct);
                if (result.MessageType == WebSocketMessageType.Close)
                    throw new InvalidOperationException("WebSocket closed unexpectedly while reading");

                var chunk = Encoding.ASCII.GetString(_recvBuffer, 0, result.Count);
                _pending.Append(chunk);
            }
        }
    }
}