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
@@ -1,4 +1,5 @@
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.ExceptionServices;
 using System.Security.Cryptography;
 using Google.Protobuf.WellKnownTypes;
 using Microsoft.Extensions.Logging;
@@ -20,6 +21,12 @@ public sealed class SessionManager : ISessionManager
    public const string DetachGraceExpiredReason = "detach-grace-expired";
    public const string FaultedReason = "faulted-reaped";

+    // Bounded so a mass expiry (or a host stop with a full registry) cannot stampede
+    // worker-process teardown: every concurrent close is one x86 worker being shut down or
+    // killed, and the point of the fan-out is to hide a few hung workers, not to tear the whole
+    // registry down at once.
+    private const int MaxParallelSessionCloses = 4;
+
    private readonly ISessionRegistry _registry;
    private readonly ISessionWorkerClientFactory _workerClientFactory;
    private readonly GatewayMetrics _metrics;
@@ -252,7 +259,10 @@ public sealed class SessionManager : ISessionManager
        DateTimeOffset now,
        CancellationToken cancellationToken)
    {
-        int closedCount = 0;
+        // Selection phase — deliberately sequential. Only the close calls below run in parallel:
+        // deciding WHICH sessions to close must stay a single ordered pass so the sweep-precedence
+        // rule and the TOCTOU re-check keep their meaning.
+        List<(GatewaySession Session, string Reason)> selected = [];
        foreach (GatewaySession session in _registry.Snapshot())
        {
            // A session is swept when its normal lease has expired, it has FAULTED (a faulted
@@ -288,45 +298,100 @@ public sealed class SessionManager : ISessionManager
                continue;
            }

-            await CloseSessionCoreAsync(session, reason, cancellationToken).ConfigureAwait(false);
-            closedCount++;
+            selected.Add((session, reason));
        }

+        if (selected.Count == 0)
+        {
+            return 0;
+        }
+
+        int closedCount = 0;
+        object failureSyncRoot = new();
+        ExceptionDispatchInfo? firstFailure = null;
+
+        // Close phase. Each close is bounded by the worker shutdown timeout (default 10 s), so a
+        // mass expiry with a few hung workers would serialize reaping and starve session slots.
+        // Parallel close is safe because TryBeginCloseIfExpired above already flipped every
+        // selected session to Closing under its own lock — that idempotent begin-close is the
+        // per-session exclusivity invariant, so no two teardowns can run against one session and
+        // a session selected here cannot be re-selected by a concurrent sweep.
+        await Parallel.ForEachAsync(
+            selected,
+            new ParallelOptions
+            {
+                MaxDegreeOfParallelism = MaxParallelSessionCloses,
+                CancellationToken = cancellationToken,
+            },
+            async (candidate, closeToken) =>
+            {
+                try
+                {
+                    await CloseSessionCoreAsync(candidate.Session, candidate.Reason, closeToken).ConfigureAwait(false);
+                    Interlocked.Increment(ref closedCount);
+                }
+                catch (Exception exception)
+                {
+                    // The sequential sweep let a close failure propagate to the lease monitor,
+                    // which logs it; that signal is preserved by rethrowing the first failure
+                    // below. It is captured rather than thrown here so one failed (or hung)
+                    // teardown does not abandon the rest of the already-selected set.
+                    lock (failureSyncRoot)
+                    {
+                        firstFailure ??= ExceptionDispatchInfo.Capture(exception);
+                    }
+                }
+            }).ConfigureAwait(false);
+
+        firstFailure?.Throw();
+
        return closedCount;
    }

    /// <inheritdoc />
    public async Task ShutdownAsync(CancellationToken cancellationToken)
    {
-        foreach (GatewaySession session in _registry.Snapshot())
-        {
-            try
+        // Sessions are drained in parallel: at a worst-case worker shutdown timeout each, a
+        // one-at-a-time drain of a full registry outruns any host stop-timeout and leaves the
+        // tail to the orphan killer. Per-session exclusivity comes from GatewaySession.CloseAsync's
+        // own close gate, and each iteration touches only its own session plus thread-safe
+        // registry/metrics state.
+        await Parallel.ForEachAsync(
+            _registry.Snapshot(),
+            new ParallelOptions
            {
-                await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
-            }
-            catch (Exception exception)
+                MaxDegreeOfParallelism = MaxParallelSessionCloses,
+                CancellationToken = cancellationToken,
+            },
+            async (session, closeToken) =>
            {
-                _logger.LogWarning(
-                    exception,
-                    "Graceful shutdown failed for session {SessionId}; killing worker.",
-                    session.SessionId);
-
-                if (_registry.TryGet(session.SessionId, out _))
+                try
                {
-                    try
+                    await CloseSessionCoreAsync(session, GatewayShutdownReason, closeToken).ConfigureAwait(false);
+                }
+                catch (Exception exception)
+                {
+                    _logger.LogWarning(
+                        exception,
+                        "Graceful shutdown failed for session {SessionId}; killing worker.",
+                        session.SessionId);
+
+                    if (_registry.TryGet(session.SessionId, out _))
                    {
-                        await KillWorkerAsync(session.SessionId, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
-                    }
-                    catch (SessionManagerException killException)
-                    {
-                        _logger.LogWarning(
-                            killException,
-                            "Worker kill fallback failed for session {SessionId}.",
-                            session.SessionId);
+                        try
+                        {
+                            await KillWorkerAsync(session.SessionId, GatewayShutdownReason, closeToken).ConfigureAwait(false);
+                        }
+                        catch (SessionManagerException killException)
+                        {
+                            _logger.LogWarning(
+                                killException,
+                                "Worker kill fallback failed for session {SessionId}.",
+                                session.SessionId);
+                        }
                    }
                }
-            }
-        }
+            }).ConfigureAwait(false);
    }

    private async Task<SessionCloseResult> CloseSessionCoreAsync(
@@ -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;
+            }
+        }
+    }
 }