natsdotnet/tests/NATS.Server.Benchmark.Tests/Transport/TlsPubSubTests.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.Transport;

[Collection("Benchmark-Tls")]
public class TlsPubSubTests(TlsServerFixture fixture, ITestOutputHelper output)
{
    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task TlsPubSub1To1_128B()
    {
        const int payloadSize = 128;
        const int messageCount = 10_000;

        var dotnetResult = await RunTlsPubSub("TLS PubSub 1:1 (128B)", "DotNet", fixture.CreateDotNetTlsClient, payloadSize, messageCount);

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

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

        var dotnetResult = await RunTlsPubOnly("TLS Pub-Only (128B)", "DotNet", fixture.CreateDotNetTlsClient, payloadSize);

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

    private static async Task<BenchmarkResult> RunTlsPubSub(string name, string serverType, Func<NatsConnection> createClient, int payloadSize, int messageCount)
    {
        var payload = new byte[payloadSize];
        var subject = $"bench.tls.pubsub.{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);

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

        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;
                }
            }
        });

        var sw = System.Diagnostics.Stopwatch.StartNew();
        for (var i = 0; i < messageCount; i++)
            await pubClient.PublishAsync(subject, payload);
        await pubClient.PingAsync();

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));
        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,
        };
    }

    private static async Task<BenchmarkResult> RunTlsPubOnly(string name, string serverType, Func<NatsConnection> createClient, int payloadSize)
    {
        var subject = $"bench.tls.pubonly.{Guid.NewGuid():N}";

        await using var client = createClient();
        await client.ConnectAsync();

        var runner = new BenchmarkRunner { WarmupCount = 1_000, MeasurementCount = 100_000 };

        return await runner.MeasureThroughputAsync(
            name,
            serverType,
            payloadSize,
            async _ => await client.PublishAsync(subject, new byte[payloadSize]));
    }
}