perf(sessions): bounded-parallel teardown in lease sweep and shutdown

2026-08-15 12:43:30 -04:00
parent a1a38b5538
commit 0bc13b5292
3 changed files with 267 additions and 36 deletions
@@ -1087,6 +1087,90 @@ public sealed class SessionManagerTests
        Assert.Equal(1, workerClient.ShutdownCount);
    }

+    /// <summary>
+    ///     A sweep pass tears the selected sessions down concurrently rather than one worker
+    ///     shutdown after another: with a mass expiry, a few hung workers would otherwise
+    ///     serialize reaping (each close is bounded by <c>Worker:ShutdownTimeoutSeconds</c>) and
+    ///     starve session slots. Overlap is asserted by counting concurrent entries into the fake
+    ///     worker's shutdown rather than by wall clock, which is sturdier on a loaded box.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [Fact]
+    public async Task CloseExpiredLeasesAsync_ClosesExpiredSessionsConcurrently()
+    {
+        ShutdownConcurrencyProbe probe = new(expectedConcurrency: 2);
+        FakeWorkerClient firstClient = new() { ShutdownConcurrencyProbe = probe };
+        FakeWorkerClient secondClient = new() { ShutdownConcurrencyProbe = probe };
+        SessionManager manager = CreateManager(new QueueingSessionWorkerClientFactory(firstClient, secondClient));
+        GatewaySession firstSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
+        GatewaySession secondSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", ownerKeyId: null, CancellationToken.None);
+        DateTimeOffset now = DateTimeOffset.UtcNow;
+        firstSession.ExtendLease(now.AddSeconds(-1));
+        secondSession.ExtendLease(now.AddSeconds(-1));
+
+        int closedCount = await manager.CloseExpiredLeasesAsync(now, CancellationToken.None);
+
+        Assert.Equal(2, closedCount);
+        Assert.Equal(SessionState.Closed, firstSession.State);
+        Assert.Equal(SessionState.Closed, secondSession.State);
+        Assert.Equal(2, probe.MaxObservedConcurrency);
+    }
+
+    /// <summary>
+    ///     Host stop drains sessions concurrently: 50 sessions at a worst-case 10 s shutdown each
+    ///     would exceed any host stop-timeout if drained one at a time, leaving the tail to the
+    ///     orphan killer.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [Fact]
+    public async Task ShutdownAsync_ClosesSessionsConcurrently()
+    {
+        ShutdownConcurrencyProbe probe = new(expectedConcurrency: 2);
+        FakeWorkerClient firstClient = new() { ShutdownConcurrencyProbe = probe };
+        FakeWorkerClient secondClient = new() { ShutdownConcurrencyProbe = probe };
+        SessionManager manager = CreateManager(new QueueingSessionWorkerClientFactory(firstClient, secondClient));
+        GatewaySession firstSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
+        GatewaySession secondSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", ownerKeyId: null, CancellationToken.None);
+
+        await manager.ShutdownAsync(CancellationToken.None);
+
+        Assert.Equal(SessionState.Closed, firstSession.State);
+        Assert.Equal(SessionState.Closed, secondSession.State);
+        Assert.Equal(2, probe.MaxObservedConcurrency);
+    }
+
+    /// <summary>
+    ///     A close that throws must not abandon the rest of the selected set: the sweep still
+    ///     tears the healthy expired session down, and the failure still surfaces to the lease
+    ///     monitor (which logs it) exactly as the sequential loop did.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [Fact]
+    public async Task CloseExpiredLeasesAsync_WhenOneCloseFails_StillClosesRemainingSessionsAndRethrows()
+    {
+        FakeWorkerClient failingClient = new()
+        {
+            ShutdownException = new InvalidOperationException("worker shutdown failed"),
+            KillException = new InvalidOperationException("worker kill failed"),
+        };
+        FakeWorkerClient healthyClient = new();
+        SessionManager manager = CreateManager(new QueueingSessionWorkerClientFactory(failingClient, healthyClient));
+        GatewaySession failingSession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
+        GatewaySession healthySession = await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", ownerKeyId: null, CancellationToken.None);
+        DateTimeOffset now = DateTimeOffset.UtcNow;
+        failingSession.ExtendLease(now.AddSeconds(-1));
+        healthySession.ExtendLease(now.AddSeconds(-1));
+
+        SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
+            async () => await manager.CloseExpiredLeasesAsync(now, CancellationToken.None));
+
+        Assert.Equal(SessionManagerErrorCode.CloseFailed, exception.ErrorCode);
+        Assert.Equal(1, healthyClient.ShutdownCount);
+        Assert.Equal(SessionState.Closed, healthySession.State);
+        Assert.False(manager.TryGetSession(healthySession.SessionId, out _));
+        Assert.False(manager.TryGetSession(failingSession.SessionId, out _));
+    }
+
    /// <summary>Verifies that shutdown closes all registered sessions.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
@@ -1274,6 +1358,13 @@ public sealed class SessionManagerTests
        /// <summary>Gets a value indicating whether to block shutdown on the fake worker client.</summary>
        public bool BlockShutdown { get; init; }

+        /// <summary>
+        ///     Gets the rendezvous that records how many shutdowns overlap, shared by the fakes of
+        ///     the sessions a single teardown pass closes. Null when the test does not measure
+        ///     teardown concurrency.
+        /// </summary>
+        public ShutdownConcurrencyProbe? ShutdownConcurrencyProbe { get; init; }
+
        /// <summary>Gets the last command invoked on the fake worker client.</summary>
        public WorkerCommand? LastCommand { get; private set; }

@@ -1335,6 +1426,11 @@ public sealed class SessionManagerTests
                throw ShutdownException;
            }

+            if (ShutdownConcurrencyProbe is not null)
+            {
+                await ShutdownConcurrencyProbe.EnterAsync(cancellationToken);
+            }
+
            if (BlockShutdown)
            {
                ShutdownStarted.TrySetResult();
@@ -1379,4 +1475,65 @@ public sealed class SessionManagerTests
        }
    }

+    /// <summary>
+    ///     Rendezvous that measures how many worker shutdowns a teardown pass runs at once. Each
+    ///     entering shutdown records the in-flight count and waits until <paramref name="expectedConcurrency"/>
+    ///     shutdowns are in flight, so a genuinely parallel teardown releases immediately while a
+    ///     sequential one can only release on the bounded timeout — with a max observed concurrency
+    ///     of one, which is the assertion that fails.
+    /// </summary>
+    /// <param name="expectedConcurrency">Number of overlapping shutdowns that releases the rendezvous.</param>
+    private sealed class ShutdownConcurrencyProbe(int expectedConcurrency)
+    {
+        private static readonly TimeSpan RendezvousTimeout = TimeSpan.FromSeconds(5);
+
+        private readonly TaskCompletionSource _reached = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private int _inFlight;
+        private int _maxInFlight;
+
+        /// <summary>Gets the highest number of shutdowns observed in flight at the same time.</summary>
+        public int MaxObservedConcurrency => Volatile.Read(ref _maxInFlight);
+
+        /// <summary>Enters the rendezvous for one worker shutdown and waits for the expected overlap.</summary>
+        /// <param name="cancellationToken">Token that abandons the wait.</param>
+        /// <returns>A task that represents the asynchronous operation.</returns>
+        public async Task EnterAsync(CancellationToken cancellationToken)
+        {
+            int inFlight = Interlocked.Increment(ref _inFlight);
+            RecordMax(inFlight);
+            if (inFlight >= expectedConcurrency)
+            {
+                _reached.TrySetResult();
+            }
+
+            try
+            {
+                await _reached.Task.WaitAsync(RendezvousTimeout, cancellationToken);
+            }
+            catch (TimeoutException)
+            {
+                // Sequential teardown: the expected overlap never happens, so let the shutdown
+                // finish and let MaxObservedConcurrency report the (failing) truth.
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _inFlight);
+            }
+        }
+
+        private void RecordMax(int inFlight)
+        {
+            int observed = Volatile.Read(ref _maxInFlight);
+            while (inFlight > observed)
+            {
+                int previous = Interlocked.CompareExchange(ref _maxInFlight, inFlight, observed);
+                if (previous == observed)
+                {
+                    return;
+                }
+
+                observed = previous;
+            }
+        }
+    }
 }