natsdotnet/tests/NATS.Server.Benchmark.Tests/CorePubSub/PubSubOneToOneTests.cs

using NATS.Client.Core;
using NATS.Server.Benchmark.Tests.Harness;
using NATS.Server.Benchmark.Tests.Infrastructure;
using Xunit.Abstractions;

namespace NATS.Server.Benchmark.Tests.CorePubSub;

[Collection("Benchmark-Core")]
public class PubSubOneToOneTests(CoreServerPairFixture fixture, ITestOutputHelper output)
{
    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task PubSub1To1_16B()
    {
        const int payloadSize = 16;
        const int messageCount = 50_000;

        var dotnetResult = await RunPubSub("PubSub 1:1 (16B)", "DotNet", payloadSize, messageCount, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunPubSub("PubSub 1:1 (16B)", "Go", payloadSize, messageCount, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task PubSub1To1_16KB()
    {
        const int payloadSize = 16 * 1024;
        const int messageCount = 5_000;

        var dotnetResult = await RunPubSub("PubSub 1:1 (16KB)", "DotNet", payloadSize, messageCount, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunPubSub("PubSub 1:1 (16KB)", "Go", payloadSize, messageCount, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

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

        await using var pubClient = createClient();
        await using var subClient = createClient();
        await pubClient.ConnectAsync();
        await subClient.ConnectAsync();

        var received = 0;
        var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);

        // Start subscriber
        var sub = await subClient.SubscribeCoreAsync<byte[]>(subject);

        // Flush to ensure subscription is propagated
        await subClient.PingAsync();
        await pubClient.PingAsync();

        var subTask = Task.Run(async () =>
        {
            await foreach (var msg in sub.Msgs.ReadAllAsync())
            {
                if (Interlocked.Increment(ref received) >= messageCount)
                {
                    tcs.TrySetResult();
                    return;
                }
            }
        });

        // Publish and measure
        var sw = System.Diagnostics.Stopwatch.StartNew();
        for (var i = 0; i < messageCount; i++)
            await pubClient.PublishAsync(subject, payload);
        await pubClient.PingAsync(); // Flush all pending writes

        // Wait for all messages received
        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(120));
        await tcs.Task.WaitAsync(cts.Token);
        sw.Stop();

        await sub.UnsubscribeAsync();

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