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

using System.Net;
using System.Net.Sockets;
using Microsoft.Extensions.Logging.Abstractions;

namespace NATS.Server.Core.Tests;

public class AcceptLoopReloadLockTests
{
    [Fact]
    public async Task Accept_loop_blocks_client_creation_while_reload_lock_is_held()
    {
        await using var fx = await AcceptLoopFixture.StartAsync();
        await fx.HoldReloadLockAsync();
        (await fx.TryConnectClientAsync(timeoutMs: 150)).ShouldBeFalse();
        fx.ReleaseReloadLock();
    }
}

internal sealed class AcceptLoopFixture : IAsyncDisposable
{
    private readonly NatsServer _server;
    private readonly CancellationTokenSource _cts;
    private bool _reloadHeld;

    private AcceptLoopFixture(NatsServer server, CancellationTokenSource cts)
    {
        _server = server;
        _cts = cts;
    }

    public static async Task<AcceptLoopFixture> StartAsync()
    {
        var server = new NatsServer(new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
        }, NullLoggerFactory.Instance);
        var cts = new CancellationTokenSource();
        _ = server.StartAsync(cts.Token);
        await server.WaitForReadyAsync();
        return new AcceptLoopFixture(server, cts);
    }

    public async Task HoldReloadLockAsync()
    {
        await _server.AcquireReloadLockForTestAsync();
        _reloadHeld = true;
    }

    public void ReleaseReloadLock()
    {
        if (_reloadHeld)
        {
            _server.ReleaseReloadLockForTest();
            _reloadHeld = false;
        }
    }

    public async Task<bool> TryConnectClientAsync(int timeoutMs)
    {
        using var client = new TcpClient();
        using var timeout = new CancellationTokenSource(timeoutMs);
        await client.ConnectAsync(IPAddress.Loopback, _server.Port, timeout.Token);
        await using var stream = client.GetStream();
        var buffer = new byte[1];
        try
        {
            var read = await stream.ReadAsync(buffer.AsMemory(0, 1), timeout.Token);
            return read > 0;
        }
        catch (OperationCanceledException)
        {
            return false;
        }
    }

    public async ValueTask DisposeAsync()
    {
        ReleaseReloadLock();
        await _cts.CancelAsync();
        _server.Dispose();
        _cts.Dispose();
    }
}