natsdotnet/tests/NATS.Server.Core.Tests/RttTests.cs

using System.Net;
using System.Net.Http.Json;
using System.Net.Sockets;
using Microsoft.Extensions.Logging.Abstractions;
using NATS.Server.Monitoring;
using NATS.Server.TestUtilities;

namespace NATS.Server.Core.Tests;

public class RttTests : IAsyncLifetime
{
    private readonly NatsServer _server;
    private readonly int _natsPort;
    private readonly int _monitorPort;
    private readonly CancellationTokenSource _cts = new();
    private readonly HttpClient _http = new();

    public RttTests()
    {
        _natsPort = TestPortAllocator.GetFreePort();
        _monitorPort = TestPortAllocator.GetFreePort();
        _server = new NatsServer(
            new NatsOptions
            {
                Port = _natsPort,
                MonitorPort = _monitorPort,
                PingInterval = TimeSpan.FromMilliseconds(200),
                MaxPingsOut = 4,
            },
            NullLoggerFactory.Instance);
    }

    public async Task InitializeAsync()
    {
        _ = _server.StartAsync(_cts.Token);
        await _server.WaitForReadyAsync();
        for (int i = 0; i < 50; i++)
        {
            try
            {
                var resp = await _http.GetAsync($"http://127.0.0.1:{_monitorPort}/healthz");
                if (resp.IsSuccessStatusCode) break;
            }
            catch (HttpRequestException) { }
            await Task.Delay(50);
        }
    }

    public async Task DisposeAsync()
    {
        _http.Dispose();
        await _cts.CancelAsync();
        _server.Dispose();
    }

    [Fact]
    public async Task Rtt_populated_after_ping_pong_cycle()
    {
        using var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await sock.ConnectAsync(new IPEndPoint(IPAddress.Loopback, _natsPort));
        using var stream = new NetworkStream(sock);
        var buf = new byte[4096];
        _ = await stream.ReadAsync(buf); // INFO

        // Send CONNECT + PING (triggers firstPongSent)
        await stream.WriteAsync("CONNECT {}\r\nPING\r\n"u8.ToArray());
        await stream.FlushAsync();
        _ = await stream.ReadAsync(buf); // PONG

        // Wait for server's PING cycle
        await Task.Delay(500);

        // Read server PING and respond with PONG
        var received = new byte[4096];
        int totalRead = 0;
        bool gotPing = false;
        using var readCts = new CancellationTokenSource(2000);
        while (!gotPing && !readCts.IsCancellationRequested)
        {
            var n = await stream.ReadAsync(received.AsMemory(totalRead), readCts.Token);
            totalRead += n;
            var text = System.Text.Encoding.ASCII.GetString(received, 0, totalRead);
            if (text.Contains("PING"))
            {
                gotPing = true;
                await stream.WriteAsync("PONG\r\n"u8.ToArray());
                await stream.FlushAsync();
            }
        }

        gotPing.ShouldBeTrue("Server should have sent PING");

        // Wait for RTT to be computed
        await Task.Delay(200);

        var response = await _http.GetAsync($"http://127.0.0.1:{_monitorPort}/connz");
        var connz = await response.Content.ReadFromJsonAsync<Connz>();
        connz.ShouldNotBeNull();
        var conn = connz.Conns.FirstOrDefault(c => c.Rtt != "");
        conn.ShouldNotBeNull("At least one connection should have RTT populated");
    }

    [Fact]
    public async Task Connz_sort_by_rtt()
    {
        using var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await sock.ConnectAsync(new IPEndPoint(IPAddress.Loopback, _natsPort));
        using var stream = new NetworkStream(sock);
        var buf = new byte[4096];
        _ = await stream.ReadAsync(buf);
        await stream.WriteAsync("CONNECT {}\r\n"u8.ToArray());
        await Task.Delay(200);

        var response = await _http.GetAsync($"http://127.0.0.1:{_monitorPort}/connz?sort=rtt");
        response.StatusCode.ShouldBe(System.Net.HttpStatusCode.OK);
    }

}