natsdotnet/tests/NATS.Server.Benchmark.Tests/Mqtt/MqttThroughputTests.cs

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

namespace NATS.Server.Benchmark.Tests.Mqtt;

[Collection("Benchmark-Mqtt")]
public class MqttThroughputTests(MqttServerFixture fixture, ITestOutputHelper output)
{
    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task MqttPubSub_128B()
    {
        const int payloadSize = 128;
        const int messageCount = 5_000;

        var dotnetResult = await RunMqttPubSub("MQTT PubSub (128B)", "DotNet", fixture.DotNetMqttPort, payloadSize, messageCount);

        if (fixture.GoAvailable)
        {
            var goResult = await RunMqttPubSub("MQTT PubSub (128B)", "Go", fixture.GoMqttPort, payloadSize, messageCount);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

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

        var dotnetResult = await RunCrossProtocol("Cross-Protocol NATS→MQTT (128B)", "DotNet", fixture.DotNetMqttPort, fixture.CreateDotNetNatsClient, payloadSize, messageCount);

        if (fixture.GoAvailable)
        {
            var goResult = await RunCrossProtocol("Cross-Protocol NATS→MQTT (128B)", "Go", fixture.GoMqttPort, fixture.CreateGoNatsClient, payloadSize, messageCount);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    private static async Task<BenchmarkResult> RunMqttPubSub(string name, string serverType, int mqttPort, int payloadSize, int messageCount)
    {
        var payload = new byte[payloadSize];
        var topic = $"bench/mqtt/pubsub/{Guid.NewGuid():N}";

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

        var factory = new MqttFactory();
        using var subscriber = factory.CreateMqttClient();
        using var publisher = factory.CreateMqttClient();

        var subOpts = new MqttClientOptionsBuilder()
            .WithTcpServer("127.0.0.1", mqttPort)
            .WithClientId($"bench-sub-{Guid.NewGuid():N}")
            .WithProtocolVersion(MQTTnet.Formatter.MqttProtocolVersion.V311)
            .Build();

        var pubOpts = new MqttClientOptionsBuilder()
            .WithTcpServer("127.0.0.1", mqttPort)
            .WithClientId($"bench-pub-{Guid.NewGuid():N}")
            .WithProtocolVersion(MQTTnet.Formatter.MqttProtocolVersion.V311)
            .Build();

        await subscriber.ConnectAsync(subOpts, cts.Token);
        await publisher.ConnectAsync(pubOpts, cts.Token);

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

        subscriber.ApplicationMessageReceivedAsync += _ =>
        {
            if (Interlocked.Increment(ref received) >= messageCount)
                tcs.TrySetResult();
            return Task.CompletedTask;
        };

        await subscriber.SubscribeAsync(
            factory.CreateSubscribeOptionsBuilder()
                .WithTopicFilter(topic)
                .Build(),
            cts.Token);

        await Task.Delay(200, cts.Token);

        var sw = System.Diagnostics.Stopwatch.StartNew();
        for (var i = 0; i < messageCount; i++)
        {
            await publisher.PublishAsync(
                new MqttApplicationMessageBuilder()
                    .WithTopic(topic)
                    .WithPayload(payload)
                    .Build(),
                cts.Token);
        }

        await tcs.Task.WaitAsync(cts.Token);
        sw.Stop();

        await subscriber.DisconnectAsync(cancellationToken: cts.Token);
        await publisher.DisconnectAsync(cancellationToken: cts.Token);

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

    private static async Task<BenchmarkResult> RunCrossProtocol(string name, string serverType, int mqttPort, Func<NatsConnection> createNatsClient, int payloadSize, int messageCount)
    {
        var payload = new byte[payloadSize];
        var natsSubject = $"bench.mqtt.cross.{Guid.NewGuid():N}";
        var mqttTopic = natsSubject.Replace('.', '/');

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

        var factory = new MqttFactory();
        using var mqttSub = factory.CreateMqttClient();

        var subOpts = new MqttClientOptionsBuilder()
            .WithTcpServer("127.0.0.1", mqttPort)
            .WithClientId($"bench-cross-sub-{Guid.NewGuid():N}")
            .WithProtocolVersion(MQTTnet.Formatter.MqttProtocolVersion.V311)
            .Build();

        await mqttSub.ConnectAsync(subOpts, cts.Token);

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

        mqttSub.ApplicationMessageReceivedAsync += _ =>
        {
            if (Interlocked.Increment(ref received) >= messageCount)
                tcs.TrySetResult();
            return Task.CompletedTask;
        };

        await mqttSub.SubscribeAsync(
            factory.CreateSubscribeOptionsBuilder()
                .WithTopicFilter(mqttTopic)
                .Build(),
            cts.Token);

        await Task.Delay(200, cts.Token);

        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(natsSubject, payload, cancellationToken: cts.Token);
        await natsPub.PingAsync(cts.Token);

        await tcs.Task.WaitAsync(cts.Token);
        sw.Stop();

        await mqttSub.DisconnectAsync(cancellationToken: cts.Token);

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