feat: add graceful shutdown, accept loop backoff, and task tracking

2026-02-22 23:43:25 -05:00
parent 600c6f9e5a
commit b68f898fa0
2 changed files with 232 additions and 13 deletions
--- a/tests/NATS.Server.Tests/ServerTests.cs
+++ b/tests/NATS.Server.Tests/ServerTests.cs
@@ -547,3 +547,115 @@ public class ServerIdentityTests
        server.Dispose();
    }
 }
+
+public class GracefulShutdownTests
+{
+    private static int GetFreePort()
+    {
+        using var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+        sock.Bind(new IPEndPoint(IPAddress.Loopback, 0));
+        return ((IPEndPoint)sock.LocalEndPoint!).Port;
+    }
+
+    [Fact]
+    public async Task ShutdownAsync_disconnects_all_clients()
+    {
+        var port = GetFreePort();
+        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
+        _ = server.StartAsync(CancellationToken.None);
+        await server.WaitForReadyAsync();
+
+        // Connect 2 raw TCP clients
+        using var client1 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+        await client1.ConnectAsync(IPAddress.Loopback, port);
+        var buf = new byte[4096];
+        await client1.ReceiveAsync(buf, SocketFlags.None); // INFO
+
+        using var client2 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+        await client2.ConnectAsync(IPAddress.Loopback, port);
+        await client2.ReceiveAsync(buf, SocketFlags.None); // INFO
+
+        // Send CONNECT so both are registered
+        await client1.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
+        await client2.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
+
+        // Wait for PONG from both (confirming they are registered)
+        using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        await client1.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
+        await client2.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
+
+        server.ClientCount.ShouldBe(2);
+
+        await server.ShutdownAsync();
+
+        server.ClientCount.ShouldBe(0);
+        server.Dispose();
+    }
+
+    [Fact]
+    public async Task WaitForShutdown_blocks_until_shutdown()
+    {
+        var port = GetFreePort();
+        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
+        _ = server.StartAsync(CancellationToken.None);
+        await server.WaitForReadyAsync();
+
+        // Start WaitForShutdown in background
+        var waitTask = Task.Run(() => server.WaitForShutdown());
+
+        // Give it a moment -- it should NOT complete yet
+        await Task.Delay(200);
+        waitTask.IsCompleted.ShouldBeFalse();
+
+        // Trigger shutdown
+        await server.ShutdownAsync();
+
+        // WaitForShutdown should complete within 5 seconds
+        var completed = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(5)));
+        completed.ShouldBe(waitTask);
+
+        server.Dispose();
+    }
+
+    [Fact]
+    public async Task ShutdownAsync_is_idempotent()
+    {
+        var port = GetFreePort();
+        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
+        _ = server.StartAsync(CancellationToken.None);
+        await server.WaitForReadyAsync();
+
+        // Call ShutdownAsync 3 times -- should not throw
+        await server.ShutdownAsync();
+        await server.ShutdownAsync();
+        await server.ShutdownAsync();
+
+        server.IsShuttingDown.ShouldBeTrue();
+        server.Dispose();
+    }
+
+    [Fact]
+    public async Task Accept_loop_waits_for_active_clients()
+    {
+        var port = GetFreePort();
+        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
+        _ = server.StartAsync(CancellationToken.None);
+        await server.WaitForReadyAsync();
+
+        // Connect a client
+        using var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+        await client.ConnectAsync(IPAddress.Loopback, port);
+        var buf = new byte[4096];
+        await client.ReceiveAsync(buf, SocketFlags.None); // INFO
+        await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
+        using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        await client.ReceiveAsync(buf, SocketFlags.None, readCts.Token); // PONG
+
+        // ShutdownAsync should complete within 10 seconds (doesn't hang)
+        var shutdownTask = server.ShutdownAsync();
+        var completed = await Task.WhenAny(shutdownTask, Task.Delay(TimeSpan.FromSeconds(10)));
+        completed.ShouldBe(shutdownTask);
+
+        server.Dispose();
+    }
+}