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(