CBDDC/tests/ZB.MOM.WW.CDBBC.E2E.Benchmark.Tests/E2EThroughputBenchmarks.cs

using BenchmarkDotNet.Attributes;
using System.Net;
using System.Net.Sockets;
using ZB.MOM.WW.CBDDC.Core.Network;
using ZB.MOM.WW.CBDDC.Sample.Console;

namespace ZB.MOM.WW.CDBBC.E2E.Benchmark.Tests;

[MemoryDiagnoser]
[SimpleJob(launchCount: 1, warmupCount: 1, iterationCount: 3)]
public class E2EThroughputBenchmarks
{
    private const int BatchSize = 50;
    private BenchmarkPeerNode _nodeA = null!;
    private BenchmarkPeerNode _nodeB = null!;
    private int _sequence;

    /// <summary>
    ///     Sets up benchmark nodes and prepares the cluster.
    /// </summary>
    [GlobalSetup]
    public async Task GlobalSetupAsync()
    {
        int nodeAPort = GetAvailableTcpPort();
        int nodeBPort = GetAvailableTcpPort();
        while (nodeBPort == nodeAPort)
            nodeBPort = GetAvailableTcpPort();

        string clusterToken = Guid.NewGuid().ToString("N");

        _nodeA = BenchmarkPeerNode.Create(
            "benchmark-node-a",
            nodeAPort,
            clusterToken,
            [
                new KnownPeerConfiguration
                {
                    NodeId = "benchmark-node-b",
                    Host = "127.0.0.1",
                    Port = nodeBPort
                }
            ]);

        _nodeB = BenchmarkPeerNode.Create(
            "benchmark-node-b",
            nodeBPort,
            clusterToken,
            [
                new KnownPeerConfiguration
                {
                    NodeId = "benchmark-node-a",
                    Host = "127.0.0.1",
                    Port = nodeAPort
                }
            ]);

        await _nodeA.StartAsync();
        await _nodeB.StartAsync();

        // Allow initial network loop to settle before measurements.
        await Task.Delay(500);
    }

    /// <summary>
    ///     Handles benchmark teardown for the throughput test suite.
    /// </summary>
    [GlobalCleanup]
    public Task GlobalCleanupAsync()
    {
        // Explicit Surreal embedded disposal can race native callbacks in benchmark child processes.
        // Benchmarks run out-of-process, so process teardown is used for cleanup stability.
        return Task.CompletedTask;
    }

    /// <summary>
    ///     Measures local write throughput against a single node.
    /// </summary>
    [Benchmark(Description = "Local write throughput", OperationsPerInvoke = BatchSize)]
    public async Task LocalWriteThroughput()
    {
        IReadOnlyList<string> userIds = NextUserIds("local");
        foreach (string userId in userIds)
            await _nodeA.UpsertUserAsync(CreateUser(userId));
    }

    /// <summary>
    ///     Measures replicated write throughput across two nodes.
    /// </summary>
    [Benchmark(Description = "Cross-node replicated throughput", OperationsPerInvoke = BatchSize)]
    public async Task ReplicatedWriteThroughput()
    {
        IReadOnlyList<string> userIds = NextUserIds("replicated");
        foreach (string userId in userIds)
            await _nodeA.UpsertUserAsync(CreateUser(userId));

        await WaitForReplicationAsync(userIds, TimeSpan.FromSeconds(30));
    }

    private IReadOnlyList<string> NextUserIds(string prefix)
    {
        int start = Interlocked.Add(ref _sequence, BatchSize) - BatchSize;
        string[] ids = new string[BatchSize];
        for (var i = 0; i < BatchSize; i++)
            ids[i] = $"{prefix}-{start + i:D8}";

        return ids;
    }

    private static User CreateUser(string userId)
    {
        return new User
        {
            Id = userId,
            Name = $"user-{userId}",
            Age = 30,
            Address = new Address { City = "BenchmarkCity" }
        };
    }

    private async Task WaitForReplicationAsync(IReadOnlyList<string> userIds, TimeSpan timeout)
    {
        DateTime deadline = DateTime.UtcNow.Add(timeout);
        while (DateTime.UtcNow < deadline)
        {
            bool allPresent = true;
            foreach (string userId in userIds)
                if (!_nodeB.ContainsUser(userId))
                {
                    allPresent = false;
                    break;
                }

            if (allPresent) return;

            await Task.Delay(25);
        }

        throw new TimeoutException($"Timed out waiting for replication of {userIds.Count} users.");
    }

    private static int GetAvailableTcpPort()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        return ((IPEndPoint)listener.LocalEndpoint).Port;
    }
}