fix(sessions): refuse late worker attach after close during startup; route open rollback through removal path

A session is registered before its worker is created, so an admin Kill or
the host drain could close it during the startup window and the worker
client that arrived afterwards was stored on the dead session, MarkReady
silently refused, SessionOpened recorded and the Closed session returned
to the caller with a live x86 worker nothing would ever stop.

AttachWorkerClient now returns false under _syncRoot when the session is
Closing/Closed/Faulted or disposed; the manager runs attach + MarkReady +
SessionOpened under the per-session close gate, discards (kills and
disposes) a refused worker client and fails the open with OpenFailed.
OpenRecorded gates the closed-side accounting so a never-opened session is
not counted as closed.

The open rollback goes through RemoveSessionAsync so per-session event
metrics are dropped like every other teardown.
This commit is contained in:
Joseph Doherty
2026-10-22 06:16:53 -04:00
parent 488649ce99
commit 87f23d9712
5 changed files with 389 additions and 39 deletions
@@ -26,6 +26,7 @@ public sealed class GatewaySession
private DateTimeOffset _lastClientActivityAt;
private DateTimeOffset? _leaseExpiresAt;
private bool _closeStarted;
private bool _openRecorded;
private int _activeEventSubscriberCount;
private readonly TimeSpan _detachGrace;
private readonly TimeSpan _faultedGrace;
@@ -413,13 +414,116 @@ public sealed class GatewaySession
/// Attaches the worker client for this session.
/// </summary>
/// <param name="workerClient">Worker client to attach.</param>
public void AttachWorkerClient(IWorkerClient workerClient)
/// <returns>
/// <see langword="true"/> when the client was attached; <see langword="false"/> when the
/// session is already <see cref="SessionState.Closing"/>, <see cref="SessionState.Closed"/>
/// or <see cref="SessionState.Faulted"/>, or has been disposed — the caller owns the client
/// and must discard it.
/// </returns>
/// <remarks>
/// The session is registered before its worker is created, so an admin kill or the host
/// drain can close it while the worker is still starting (up to the startup timeout). A
/// worker client arriving after that must not be stored on the dead session: nothing would
/// ever shut it down (the close already ran with no client to stop) and the open would hand
/// a closed session back to the caller as if it had succeeded. The decision is taken under
/// <c>_syncRoot</c> so it cannot interleave with the close's state write.
/// </remarks>
public bool AttachWorkerClient(IWorkerClient workerClient)
{
ArgumentNullException.ThrowIfNull(workerClient);
lock (_syncRoot)
{
if (_state is SessionState.Closing or SessionState.Closed or SessionState.Faulted
|| _eventDistributorDisposed)
{
return false;
}
_workerClient = workerClient;
return true;
}
}
/// <summary>
/// Runs the end of the open sequence — attach the worker client, mark the session
/// <see cref="SessionState.Ready"/>, record it as opened — under the per-session close
/// gate, so a concurrent <see cref="CloseAsync"/> or
/// <see cref="KillWorkerWithCloseGateAsync"/> runs either entirely before it (the attach
/// is refused) or entirely after it (an ordinary close of a Ready session). Without the
/// gate a kill could land between the attach and the open accounting and the two teardown
/// bookkeepings would disagree about whether the session ever counted as open.
/// </summary>
/// <param name="workerClient">Worker client to attach.</param>
/// <param name="onReady">
/// Invoked inside the gate once the session is Ready; the caller records its open-session
/// accounting here, and <see cref="OpenRecorded"/> is set right after it returns.
/// </param>
/// <param name="cancellationToken">Token to cancel waiting for the gate.</param>
/// <returns>
/// <see langword="true"/> when the session became Ready; <see langword="false"/> when the
/// attach was refused (see <see cref="AttachWorkerClient"/>), including when the session was
/// disposed before the gate could be taken.
/// </returns>
internal async ValueTask<bool> TryAttachWorkerClientAndMarkReadyAsync(
IWorkerClient workerClient,
Action onReady,
CancellationToken cancellationToken)
{
ArgumentNullException.ThrowIfNull(workerClient);
ArgumentNullException.ThrowIfNull(onReady);
try
{
await _closeLock.WaitAsync(cancellationToken).ConfigureAwait(false);
}
catch (ObjectDisposedException)
{
// Disposed while the worker was starting: the session has been torn down.
return false;
}
try
{
if (!AttachWorkerClient(workerClient))
{
return false;
}
MarkReady();
onReady();
lock (_syncRoot)
{
_openRecorded = true;
}
return true;
}
finally
{
_closeLock.Release();
}
}
/// <summary>
/// Gets a value indicating whether the open-session accounting
/// (<c>GatewayMetrics.SessionOpened</c>) has been recorded for this session. Every
/// teardown path consults it so a session closed or killed while its worker was still
/// starting — never counted as opened — is not counted as closed either.
/// </summary>
/// <remarks>
/// Written only by <see cref="TryAttachWorkerClientAndMarkReadyAsync"/>, inside the close
/// gate; read by teardown paths that hold (or just released) the same gate, so the two
/// accountings never observe a half-finished open.
/// </remarks>
internal bool OpenRecorded
{
get
{
lock (_syncRoot)
{
return _openRecorded;
}
}
}
@@ -107,7 +107,7 @@ public sealed class SessionManager : ISessionManager
EnsureSessionCapacity();
GatewaySession? session = null;
bool sessionOpenedRecorded = false;
bool registered = false;
try
{
session = CreateSession(request, clientIdentity, ownerKeyId, ownerDashboardTags);
@@ -118,36 +118,31 @@ public sealed class SessionManager : ISessionManager
$"Session id collision while opening session {session.SessionId}.");
}
registered = true;
session.TransitionTo(SessionState.StartingWorker);
IWorkerClient workerClient = await _workerClientFactory
.CreateAsync(session, cancellationToken)
.ConfigureAwait(false);
session.AttachWorkerClient(workerClient);
session.MarkReady();
_metrics.SessionOpened();
sessionOpenedRecorded = true;
// The session has been registered (and killable / drainable) for the whole worker
// startup. If a kill or the host drain closed it meanwhile, the attach is refused and
// the late worker is discarded here — it must not be returned on a dead session, and
// nothing else would ever stop it.
bool ready = await session
.TryAttachWorkerClientAndMarkReadyAsync(workerClient, _metrics.SessionOpened, cancellationToken)
.ConfigureAwait(false);
if (!ready)
{
await DiscardWorkerClientAsync(session, workerClient).ConfigureAwait(false);
throw new InvalidOperationException(
$"Session {session.SessionId} was closed while its worker was starting.");
}
return session;
}
catch (Exception exception)
{
session?.MarkFaulted(exception.Message);
if (session is not null)
{
_registry.TryRemove(session.SessionId, out _);
await session.DisposeAsync().ConfigureAwait(false);
}
// If SessionOpened() already incremented the open-session gauge,
// a failure after that point (e.g. auto-subscribe rejection) must
// decrement it again so mxgateway.sessions.open does not leak.
if (sessionOpenedRecorded)
{
_metrics.SessionRemoved();
}
ReleaseSessionSlot();
await RollBackOpenAsync(session, registered, exception).ConfigureAwait(false);
_metrics.Fault(SessionManagerErrorCode.OpenFailed.ToString());
_logger.LogWarning(
exception,
@@ -161,6 +156,62 @@ public sealed class SessionManager : ISessionManager
}
}
// Order — fault, deregister (dropping per-session event metrics), release the slot, dispose —
// matters: releasing the slot before disposal would let the next open race this worker's
// teardown on the same machine. A registered session goes through RemoveSessionAsync, the
// single removal path, which is a no-op when a concurrent kill or drain already removed it
// (that teardown released the slot and disposed the session). The open-session gauge is never
// touched here: SessionOpened is recorded inside the attach gate, after which nothing in the
// open sequence can fail.
private async Task RollBackOpenAsync(GatewaySession? session, bool registered, Exception cause)
{
if (session is null)
{
ReleaseSessionSlot();
return;
}
session.MarkFaulted(cause.Message);
if (registered)
{
await RemoveSessionAsync(session).ConfigureAwait(false);
return;
}
ReleaseSessionSlot();
await session.DisposeAsync().ConfigureAwait(false);
}
// Kills and disposes a worker client that arrived after its session was closed. Best-effort:
// the client is already orphaned from the session, so a failure here is logged rather than
// allowed to mask the open failure the caller is about to receive.
private async Task DiscardWorkerClientAsync(GatewaySession session, IWorkerClient workerClient)
{
try
{
workerClient.Kill("session-closed-during-startup");
}
catch (Exception exception)
{
_logger.LogWarning(
exception,
"Failed to kill the late worker for closed session {SessionId}.",
session.SessionId);
}
try
{
await workerClient.DisposeAsync().ConfigureAwait(false);
}
catch (Exception exception)
{
_logger.LogWarning(
exception,
"Failed to dispose the late worker client for closed session {SessionId}.",
session.SessionId);
}
}
/// <inheritdoc />
public bool TryGetSession(
string sessionId,
@@ -231,11 +282,14 @@ public sealed class SessionManager : ISessionManager
session.MarkFaulted(exception.Message);
_metrics.Fault(SessionManagerErrorCode.CloseFailed.ToString());
// The open-session gauge was incremented in OpenSessionAsync; every
// session reaching KillWorkerAsync had SessionOpened recorded. If the
// kill path throws, decrement the gauge here so mxgateway.sessions.open
// does not leak — mirroring the equivalent guard in OpenSessionAsync.
_metrics.SessionRemoved();
// If the open-session gauge was incremented for this session (it was not when the
// kill landed while the worker was still starting), decrement it here so
// mxgateway.sessions.open does not leak when the kill path throws.
if (session.OpenRecorded)
{
_metrics.SessionRemoved();
}
await RemoveSessionAsync(session).ConfigureAwait(false);
throw new SessionManagerException(
SessionManagerErrorCode.CloseFailed,
@@ -243,7 +297,7 @@ public sealed class SessionManager : ISessionManager
exception);
}
if (!wasClosed)
if (!wasClosed && session.OpenRecorded)
{
_metrics.SessionClosed();
}
@@ -435,7 +489,7 @@ public sealed class SessionManager : ISessionManager
try
{
SessionCloseResult result = await session.CloseAsync(reason, cancellationToken).ConfigureAwait(false);
if (!wasClosed && !result.AlreadyClosed)
if (!wasClosed && !result.AlreadyClosed && session.OpenRecorded)
{
_metrics.SessionClosed();
}
@@ -450,7 +504,7 @@ public sealed class SessionManager : ISessionManager
catch (SessionCloseStartedException exception)
{
session.MarkFaulted(exception.Message);
if (!wasClosed)
if (!wasClosed && session.OpenRecorded)
{
// Account the close as a SessionClosed (decrements the open-session
// gauge AND increments the sessions.closed counter), not just