fix(sessions): refuse OpenSession while the gateway is stopping; drain until the registry is empty

ShutdownAsync iterated a single registry snapshot while Kestrel still
accepted OpenSession calls (ApplicationStopping fires before hosted
services stop), so a session opened behind the snapshot had nobody left to
close it and its x86 worker outlived the gateway as an orphan.

SessionManager now observes IHostApplicationLifetime.ApplicationStopping
(and sets the same flag at the start of ShutdownAsync) and rejects
OpenSessionAsync with the new SessionManagerErrorCode.GatewayStopping,
re-checked after registration so the flag write before the drain's
snapshot closes the race. GatewayStopping maps to gRPC Unavailable and is
not a fault. ShutdownAsync re-snapshots until the registry is empty,
bounded at three passes.
This commit is contained in:
Joseph Doherty
2026-10-22 06:23:40 -04:00
parent 09666cc719
commit 37ab629d70
6 changed files with 317 additions and 62 deletions
+2
View File
@@ -306,6 +306,8 @@ StatusCode statusCode = exception.ErrorCode switch
SessionManagerErrorCode.SessionLimitExceeded => StatusCode.ResourceExhausted,
SessionManagerErrorCode.OpenFailed => StatusCode.Unavailable,
SessionManagerErrorCode.CloseFailed => StatusCode.Unavailable,
SessionManagerErrorCode.PermissionDenied => StatusCode.PermissionDenied,
SessionManagerErrorCode.GatewayStopping => StatusCode.Unavailable,
_ => StatusCode.Unavailable,
};
```
+41 -24
View File
@@ -164,6 +164,8 @@ public sealed record SessionCloseResult(
| `SessionLimitExceeded` | `MaxSessions` is in use. |
| `OpenFailed` | `OpenSessionAsync` failed; the inner exception carries the cause. |
| `CloseFailed` | A close started but did not complete cleanly; the session is removed and faulted. |
| `PermissionDenied` | The caller may not perform the operation on this session (for example, attaching `StreamEvents` to a session opened by a different API key). |
| `GatewayStopping` | The gateway is stopping and refuses new sessions (see [Shutdown Coordination](#shutdown-coordination)). Maps to gRPC `Unavailable`; not recorded as a fault. |
## Lifecycle
@@ -306,40 +308,55 @@ If both graceful shutdown and the kill fall-back fail, the original and kill exc
`SessionShutdownHostedService.StopAsync` calls `SessionManager.ShutdownAsync`, which closes every registered session with `GatewayShutdownReason`. Sessions are drained with the same bounded fan-out the lease sweep uses (`MaxParallelSessionCloses`), because a one-at-a-time drain of a full registry at a worst-case worker shutdown timeout each outruns any host stop-timeout and leaves the tail to the orphan killer. Each iteration catches its own exceptions — *every* exception, including from the fallback — calls `KillWorkerAsync` on an uncancellable token, and removes the session, so that neither one stuck worker nor one failing teardown can block or abort the rest of the host's drain:
```csharp
await Parallel.ForEachAsync(
_registry.Snapshot(),
new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses },
async (session, _) =>
MarkStopping();
for (int pass = 0; pass < MaxShutdownDrainPasses; pass++)
{
if (_registry.Count == 0)
{
try
return;
}
await Parallel.ForEachAsync(
_registry.Snapshot(),
new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses },
async (session, _) =>
{
await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
}
catch (Exception exception)
{
_logger.LogWarning(
exception,
"Graceful shutdown failed for session {SessionId}; killing worker.",
session.SessionId);
if (_registry.TryGet(session.SessionId, out GatewaySession? registeredSession)
&& registeredSession is not null)
try
{
try
await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
}
catch (Exception exception)
{
_logger.LogWarning(
exception,
"Graceful shutdown failed for session {SessionId}; killing worker.",
session.SessionId);
if (_registry.TryGet(session.SessionId, out GatewaySession? registeredSession)
&& registeredSession is not null)
{
// Not the caller's token: the kill is the last-resort orphan preventer.
await KillWorkerAsync(session.SessionId, GatewayShutdownReason, CancellationToken.None).ConfigureAwait(false);
}
catch (Exception killException)
{
_logger.LogWarning(killException, "Worker kill fallback failed for session {SessionId}.", session.SessionId);
try
{
// Not the caller's token: the kill is the last-resort orphan preventer.
await KillWorkerAsync(session.SessionId, GatewayShutdownReason, CancellationToken.None).ConfigureAwait(false);
}
catch (Exception killException)
{
_logger.LogWarning(killException, "Worker kill fallback failed for session {SessionId}.", session.SessionId);
}
}
}
}
}).ConfigureAwait(false);
}).ConfigureAwait(false);
}
```
Iterating over `Snapshot` rather than the live dictionary lets `RemoveSessionAsync` mutate the registry from inside the loop without throwing, and gives the parallel drain a stable, already-materialized source.
### Stopping gate
A drain that only ever looked at one snapshot could be outrun by `OpenSession`: the host raises `ApplicationStopping` *before* hosted services stop, and Kestrel keeps accepting gRPC calls until its own `StopAsync` runs, so a session opened behind the snapshot would have nobody left to close it and its x86 worker would survive the gateway as an orphan. `SessionManager` therefore keeps a stopping flag that is set by the host's `IHostApplicationLifetime.ApplicationStopping` (the manager registers for it in its constructor when the lifetime is injected) and, as a fallback for hosts without one, at the start of `ShutdownAsync`. Once set it is never cleared. `OpenSessionAsync` checks it twice: on entry, before a slot is taken, and again right after the session is registered. The second check is what closes the race — the flag is written before the drain takes its registry snapshot, so an open that does not observe it there registered its session before that snapshot and is drained, while one that does observe it rolls its own session back. Either way the caller receives `SessionManagerException` with `GatewayStopping`, which the gRPC layer maps to `Unavailable` so the client retries against the restarted gateway; it is not an open failure and records no fault metric.
The drain loop above is the backstop to that gate: it re-snapshots until the registry is empty, bounded at `MaxShutdownDrainPasses` (`3`) so a session that somehow cannot be removed cannot hold the host stop forever. Anything still registered after the last pass is logged and left to the orphan terminator on the next start.
## Dependency Injection
`SessionServiceCollectionExtensions.AddGatewaySessions` registers the four singletons and the hosted service:
@@ -1035,6 +1035,7 @@ public sealed class MxAccessGatewayService(
SessionManagerErrorCode.OpenFailed => StatusCode.Unavailable,
SessionManagerErrorCode.CloseFailed => StatusCode.Unavailable,
SessionManagerErrorCode.PermissionDenied => StatusCode.PermissionDenied,
SessionManagerErrorCode.GatewayStopping => StatusCode.Unavailable,
_ => StatusCode.Unavailable,
};
@@ -2,6 +2,7 @@ using System.Diagnostics.CodeAnalysis;
using System.Runtime.ExceptionServices;
using System.Security.Cryptography;
using Google.Protobuf.WellKnownTypes;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
@@ -27,6 +28,13 @@ public sealed class SessionManager : ISessionManager
// registry down at once.
private const int MaxParallelSessionCloses = 4;
// ShutdownAsync re-snapshots the registry until it is empty so an open that registered its
// session just after a snapshot is still drained. The stopping gate in OpenSessionAsync is
// what actually closes that window (it is re-checked after registration), so the loop is a
// backstop and a small bound is enough; anything still registered after the last pass is
// logged and left to the orphan terminator on the next start.
private const int MaxShutdownDrainPasses = 3;
private readonly ISessionRegistry _registry;
private readonly ISessionWorkerClientFactory _workerClientFactory;
private readonly GatewayMetrics _metrics;
@@ -39,6 +47,12 @@ public sealed class SessionManager : ISessionManager
private readonly Dashboard.Hubs.IDashboardEventBroadcaster? _dashboardEventBroadcaster;
private readonly ArrayAddressNormalizer? _addressNormalizer;
// 0 = accepting opens, 1 = stopping. Set by the host's ApplicationStopping signal (which fires
// before hosted services stop, while Kestrel still accepts gRPC calls) and, as a fallback for
// hosts without a lifetime, at the start of ShutdownAsync. Never reset: a manager that has
// begun draining does not come back.
private int _stopping;
/// <summary>
/// Initializes a new instance of <see cref="SessionManager"/>.
/// </summary>
@@ -61,6 +75,12 @@ public sealed class SessionManager : ISessionManager
/// metadata; handed to each session so the normalization runs at the outbound choke point.
/// Null in unit tests that do not exercise array-write ergonomics.
/// </param>
/// <param name="applicationLifetime">
/// Host lifetime whose <c>ApplicationStopping</c> closes the open gate as soon as the host
/// begins stopping — before hosted services stop and while Kestrel is still accepting gRPC
/// calls — so no session is opened behind the shutdown drain. Null in unit tests; the gate
/// then closes when <see cref="ShutdownAsync"/> starts.
/// </param>
public SessionManager(
ISessionRegistry registry,
ISessionWorkerClientFactory workerClientFactory,
@@ -71,7 +91,8 @@ public sealed class SessionManager : ISessionManager
Grpc.MxAccessGrpcMapper? eventMapper = null,
ILogger<SessionEventDistributor>? distributorLogger = null,
Dashboard.Hubs.IDashboardEventBroadcaster? dashboardEventBroadcaster = null,
ArrayAddressNormalizer? addressNormalizer = null)
ArrayAddressNormalizer? addressNormalizer = null,
IHostApplicationLifetime? applicationLifetime = null)
{
_registry = registry ?? throw new ArgumentNullException(nameof(registry));
_workerClientFactory = workerClientFactory ?? throw new ArgumentNullException(nameof(workerClientFactory));
@@ -85,6 +106,10 @@ public sealed class SessionManager : ISessionManager
_addressNormalizer = addressNormalizer;
_options = options.Value;
_sessionSlots = new SemaphoreSlim(_options.Sessions.MaxSessions, _options.Sessions.MaxSessions);
// The registration is never disposed: the manager is a host singleton that lives exactly
// as long as the lifetime it observes.
applicationLifetime?.ApplicationStopping.Register(MarkStopping);
}
/// <inheritdoc />
@@ -104,6 +129,7 @@ public sealed class SessionManager : ISessionManager
CancellationToken cancellationToken)
{
ArgumentNullException.ThrowIfNull(request);
ThrowIfStopping();
EnsureSessionCapacity();
GatewaySession? session = null;
@@ -119,6 +145,12 @@ public sealed class SessionManager : ISessionManager
}
registered = true;
// Re-check after registration. The stopping flag is written before ShutdownAsync takes
// its registry snapshot, so an open that does not observe it here registered its
// session before that snapshot and will be drained; one that does observe it rolls
// back itself. Either way no session is created behind the drain.
ThrowIfStopping();
session.TransitionTo(SessionState.StartingWorker);
IWorkerClient workerClient = await _workerClientFactory
.CreateAsync(session, cancellationToken)
@@ -140,6 +172,13 @@ public sealed class SessionManager : ISessionManager
return session;
}
catch (SessionManagerException exception) when (exception.ErrorCode == SessionManagerErrorCode.GatewayStopping)
{
// Not an open failure: the gateway is going away and the caller should retry against
// the restarted one, so no fault is recorded and the code is surfaced as-is.
await RollBackOpenAsync(session, registered, exception).ConfigureAwait(false);
throw;
}
catch (Exception exception)
{
await RollBackOpenAsync(session, registered, exception).ConfigureAwait(false);
@@ -435,49 +474,86 @@ public sealed class SessionManager : ISessionManager
// intentional: that sweep is periodic maintenance whose missed sessions are picked up by
// the next pass and, ultimately, by this drain. This drain is terminal — nothing runs after
// it — so it must not be abandoned partway.
await Parallel.ForEachAsync(
_registry.Snapshot(),
new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses },
async (session, _) =>
//
// Close the open gate BEFORE the first snapshot (see the re-check in OpenSessionAsync), then
// re-snapshot until the registry is empty so an open that registered just before the gate
// closed is drained too.
MarkStopping();
for (int pass = 0; pass < MaxShutdownDrainPasses; pass++)
{
if (_registry.Count == 0)
{
try
{
await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
}
catch (Exception exception)
{
_logger.LogWarning(
exception,
"Graceful shutdown failed for session {SessionId}; killing worker.",
session.SessionId);
return;
}
// Defensive fallback: CloseSessionCoreAsync's inner
// SessionCloseStartedException catch normally removes the session
// and accounts the close. This outer fallback only fires for
// sessions still in the registry — route through KillWorkerAsync
// so the bookkeeping is identical to the dashboard kill path.
if (_registry.TryGet(session.SessionId, out GatewaySession? registeredSession)
&& registeredSession is not null)
await Parallel.ForEachAsync(
_registry.Snapshot(),
new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses },
async (session, _) =>
{
try
{
try
await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
}
catch (Exception exception)
{
_logger.LogWarning(
exception,
"Graceful shutdown failed for session {SessionId}; killing worker.",
session.SessionId);
// Defensive fallback: CloseSessionCoreAsync's inner
// SessionCloseStartedException catch normally removes the session
// and accounts the close. This outer fallback only fires for
// sessions still in the registry — route through KillWorkerAsync
// so the bookkeeping is identical to the dashboard kill path.
if (_registry.TryGet(session.SessionId, out GatewaySession? registeredSession)
&& registeredSession is not null)
{
// Deliberately NOT the caller's token: the kill is the last-resort orphan
// preventer, so it must still run when the host stop deadline (or a
// sibling body's failure) has already cancelled the drain. It is a
// synchronous Kill plus registry/dispose bookkeeping, not a wait on
// the worker, so it cannot extend the drain meaningfully.
await KillWorkerAsync(session.SessionId, GatewayShutdownReason, CancellationToken.None).ConfigureAwait(false);
}
catch (Exception killException)
{
_logger.LogWarning(
killException,
"Worker kill fallback failed for session {SessionId}.",
session.SessionId);
try
{
// Deliberately NOT the caller's token: the kill is the last-resort orphan
// preventer, so it must still run when the host stop deadline (or a
// sibling body's failure) has already cancelled the drain. It is a
// synchronous Kill plus registry/dispose bookkeeping, not a wait on
// the worker, so it cannot extend the drain meaningfully.
await KillWorkerAsync(session.SessionId, GatewayShutdownReason, CancellationToken.None).ConfigureAwait(false);
}
catch (Exception killException)
{
_logger.LogWarning(
killException,
"Worker kill fallback failed for session {SessionId}.",
session.SessionId);
}
}
}
}
}).ConfigureAwait(false);
}).ConfigureAwait(false);
}
int remaining = _registry.Count;
if (remaining > 0)
{
_logger.LogWarning(
"Session drain finished with {Remaining} session(s) still registered after {Passes} passes; their workers are left to the orphan terminator on the next start.",
remaining,
MaxShutdownDrainPasses);
}
}
private void MarkStopping()
{
Volatile.Write(ref _stopping, 1);
}
private void ThrowIfStopping()
{
if (Volatile.Read(ref _stopping) != 0)
{
throw new SessionManagerException(
SessionManagerErrorCode.GatewayStopping,
"The gateway is stopping and is not accepting new sessions.");
}
}
private async Task<SessionCloseResult> CloseSessionCoreAsync(
@@ -17,4 +17,12 @@ public enum SessionManagerErrorCode
/// Maps to gRPC <c>PermissionDenied</c>.
/// </summary>
PermissionDenied,
/// <summary>
/// The gateway is stopping, so no new session can be opened: a session created behind the
/// shutdown drain's snapshot would have nobody left to close it and its worker would
/// outlive the gateway as an orphan. Maps to gRPC <c>Unavailable</c> so the client retries
/// against the restarted gateway.
/// </summary>
GatewayStopping,
}
@@ -1419,6 +1419,84 @@ public sealed class SessionManagerTests
Assert.False(manager.TryGetSession(secondSession.SessionId, out _));
}
/// <summary>
/// Once the drain has started, a new open must be refused with <c>GatewayStopping</c>
/// rather than accepted: a session created behind the drain's snapshot has nobody left to
/// close it, and its worker would outlive the gateway as an orphan. The refusal is not an
/// open failure — no slot is taken, nothing is registered and no fault is recorded.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task OpenSessionAsync_AfterShutdownStarted_ThrowsGatewayStopping()
{
FakeWorkerClient workerClient = new();
SessionRegistry registry = new();
using GatewayMetrics metrics = new();
SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient), registry, metrics);
await manager.ShutdownAsync(CancellationToken.None);
SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
async () => await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None));
Assert.Equal(SessionManagerErrorCode.GatewayStopping, exception.ErrorCode);
Assert.Equal(0, registry.Count);
Assert.Equal(0, metrics.GetSnapshot().SessionsOpened);
Assert.Equal(0, metrics.GetSnapshot().Faults);
}
/// <summary>
/// The host signals <c>ApplicationStopping</c> before hosted services stop, and Kestrel keeps
/// accepting gRPC calls until its own <c>StopAsync</c> runs — so the gate must close on the
/// lifetime signal, not only when <c>ShutdownAsync</c> is finally called.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task OpenSessionAsync_WhenApplicationStopping_ThrowsGatewayStopping()
{
FakeWorkerClient workerClient = new();
using FakeHostApplicationLifetime lifetime = new();
SessionManager manager = new(
new SessionRegistry(),
new FakeSessionWorkerClientFactory(workerClient),
Options.Create(CreateOptions()),
new GatewayMetrics(),
applicationLifetime: lifetime);
GatewaySession beforeStopping = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
Assert.Equal(SessionState.Ready, beforeStopping.State);
lifetime.SignalStopping();
SessionManagerException exception = await Assert.ThrowsAsync<SessionManagerException>(
async () => await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", ownerKeyId: null, CancellationToken.None));
Assert.Equal(SessionManagerErrorCode.GatewayStopping, exception.ErrorCode);
}
/// <summary>
/// An open that passed the stopping gate but registered its session after the drain took
/// its snapshot is still caught: the drain re-snapshots until the registry is empty, so the
/// late session is closed by the drain rather than left running.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task ShutdownAsync_DrainsSessionRegisteredAfterFirstSnapshot()
{
FakeWorkerClient workerClient = new();
LateFirstSnapshotRegistry registry = new();
using GatewayMetrics metrics = new();
SessionManager manager = CreateManager(new FakeSessionWorkerClientFactory(workerClient), registry, metrics);
GatewaySession session = await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
await manager.ShutdownAsync(CancellationToken.None);
Assert.True(registry.SnapshotCalls >= 2);
Assert.Equal(SessionState.Closed, session.State);
Assert.Equal(1, workerClient.ShutdownCount);
Assert.Equal(0, registry.Count);
Assert.Equal(0, metrics.GetSnapshot().OpenSessions);
Assert.Equal(1, metrics.GetSnapshot().SessionsClosed);
}
/// <summary>Verifies that shutdown closes all registered sessions.</summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
@@ -1634,6 +1712,79 @@ public sealed class SessionManagerTests
}
}
/// <summary>
/// Lifetime fake whose <see cref="ApplicationStopping"/> the test fires on demand; the
/// other two tokens never fire.
/// </summary>
private sealed class FakeHostApplicationLifetime : Microsoft.Extensions.Hosting.IHostApplicationLifetime, IDisposable
{
private readonly CancellationTokenSource _stopping = new();
/// <inheritdoc />
public CancellationToken ApplicationStarted => CancellationToken.None;
/// <inheritdoc />
public CancellationToken ApplicationStopping => _stopping.Token;
/// <inheritdoc />
public CancellationToken ApplicationStopped => CancellationToken.None;
/// <summary>Fires <see cref="ApplicationStopping"/>.</summary>
public void SignalStopping()
{
_stopping.Cancel();
}
/// <inheritdoc />
public void StopApplication()
{
SignalStopping();
}
/// <inheritdoc />
public void Dispose()
{
_stopping.Dispose();
}
}
/// <summary>
/// Registry whose first <see cref="Snapshot"/> is empty regardless of its contents,
/// simulating a session whose <c>TryAdd</c> landed just after the drain's first snapshot;
/// every later snapshot is the real one.
/// </summary>
private sealed class LateFirstSnapshotRegistry : ISessionRegistry
{
private readonly SessionRegistry _inner = new();
/// <summary>Gets how many snapshots have been taken.</summary>
public int SnapshotCalls { get; private set; }
/// <inheritdoc />
public int Count => _inner.Count;
/// <inheritdoc />
public int ActiveCount => _inner.ActiveCount;
/// <inheritdoc />
public bool TryAdd(GatewaySession session) => _inner.TryAdd(session);
/// <inheritdoc />
public bool TryGet(string sessionId, [System.Diagnostics.CodeAnalysis.MaybeNullWhen(false)] out GatewaySession session)
=> _inner.TryGet(sessionId, out session);
/// <inheritdoc />
public bool TryRemove(string sessionId, [System.Diagnostics.CodeAnalysis.MaybeNullWhen(false)] out GatewaySession session)
=> _inner.TryRemove(sessionId, out session);
/// <inheritdoc />
public IReadOnlyCollection<GatewaySession> Snapshot()
{
SnapshotCalls++;
return SnapshotCalls == 1 ? [] : _inner.Snapshot();
}
}
private sealed class FakeWorkerClient : IWorkerClient
{
private WorkerClientState _state = WorkerClientState.Ready;