diff --git a/docs/Sessions.md b/docs/Sessions.md
index 7bcbc51..5d52e87 100644
--- a/docs/Sessions.md
+++ b/docs/Sessions.md
@@ -174,14 +174,7 @@ public sealed record SessionCloseResult(
```csharp
catch (Exception exception)
{
- session?.MarkFaulted(exception.Message);
- if (session is not null)
- {
- _registry.TryRemove(session.SessionId, out _);
- await session.DisposeAsync().ConfigureAwait(false);
- }
-
- ReleaseSessionSlot();
+ await RollBackOpenAsync(session, registered, exception).ConfigureAwait(false);
_metrics.Fault(SessionManagerErrorCode.OpenFailed.ToString());
_logger.LogWarning(
exception,
@@ -195,7 +188,15 @@ catch (Exception exception)
}
```
-The order — fault, deregister, dispose, release slot, record metric, log, rethrow — matters because releasing the semaphore before disposal would let the next open race the worker process tear-down on the same machine.
+`RollBackOpenAsync` faults the session and, when it had been registered, routes it through `RemoveSessionAsync` — the single removal path every teardown uses, which deregisters, drops the session's per-session event metrics, releases the slot and disposes, in that order. The order matters because releasing the semaphore before disposal would let the next open race the worker process tear-down on the same machine. A session that was never registered (id collision) releases the slot and disposes directly. The open-session gauge is not touched by the rollback: `SessionOpened` is recorded inside the attach gate described next, after which nothing in the open sequence can fail.
+
+#### Close or kill during worker startup
+
+The session is registered — and therefore visible to the dashboard's Kill button, the lease sweep and the host drain — for the whole worker startup, which can last up to `MxGateway:Worker:StartupTimeoutSeconds`. A kill or drain that lands in that window finds no worker client to stop, flips the session to `Closed`, removes it and disposes it. The worker client that `SessionWorkerClientFactory.CreateAsync` then returns must not be stored on that dead session and handed back to the caller as a successful open: nothing would ever shut the worker down (the close already ran) and the client would hold a `Closed` session.
+
+`GatewaySession.AttachWorkerClient` therefore returns `false`, under `_syncRoot`, when the session is already `Closing`, `Closed` or `Faulted` or has been disposed. The manager runs the end of the open sequence — attach, `MarkReady`, `SessionOpened` — through `TryAttachWorkerClientAndMarkReadyAsync`, which holds the per-session close gate (`_closeLock`) so a concurrent `CloseAsync` / `KillWorkerWithCloseGateAsync` runs either entirely before it (the attach is refused) or entirely after it (an ordinary close of a Ready session). On a refused attach the manager kills and disposes the late worker client and the open fails with `OpenFailed`; the registry is already empty and the slot was released by whichever teardown closed the session.
+
+Because such a session was never counted as opened, it must not be counted as closed either. `GatewaySession.OpenRecorded` is set inside the same gate right after `SessionOpened` is recorded, and every teardown path (`KillWorkerAsync`, `CloseSessionCoreAsync`, the close-failed branch) records `SessionClosed` / `SessionRemoved` only when it is set — so `sessions.opened`, `sessions.closed` and the `sessions.open` gauge stay consistent across the startup window.
### Run
diff --git a/src/ZB.MOM.WW.MxGateway.Server/Sessions/GatewaySession.cs b/src/ZB.MOM.WW.MxGateway.Server/Sessions/GatewaySession.cs
index 1de4054..57393b4 100644
--- a/src/ZB.MOM.WW.MxGateway.Server/Sessions/GatewaySession.cs
+++ b/src/ZB.MOM.WW.MxGateway.Server/Sessions/GatewaySession.cs
@@ -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.
///
/// Worker client to attach.
- public void AttachWorkerClient(IWorkerClient workerClient)
+ ///
+ /// when the client was attached; when the
+ /// session is already ,
+ /// or , or has been disposed — the caller owns the client
+ /// and must discard it.
+ ///
+ ///
+ /// 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
+ /// _syncRoot so it cannot interleave with the close's state write.
+ ///
+ 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;
+ }
+ }
+
+ ///
+ /// Runs the end of the open sequence — attach the worker client, mark the session
+ /// , record it as opened — under the per-session close
+ /// gate, so a concurrent or
+ /// 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.
+ ///
+ /// Worker client to attach.
+ ///
+ /// Invoked inside the gate once the session is Ready; the caller records its open-session
+ /// accounting here, and is set right after it returns.
+ ///
+ /// Token to cancel waiting for the gate.
+ ///
+ /// when the session became Ready; when the
+ /// attach was refused (see ), including when the session was
+ /// disposed before the gate could be taken.
+ ///
+ internal async ValueTask 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();
+ }
+ }
+
+ ///
+ /// Gets a value indicating whether the open-session accounting
+ /// (GatewayMetrics.SessionOpened) 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.
+ ///
+ ///
+ /// Written only by , 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.
+ ///
+ internal bool OpenRecorded
+ {
+ get
+ {
+ lock (_syncRoot)
+ {
+ return _openRecorded;
+ }
}
}
diff --git a/src/ZB.MOM.WW.MxGateway.Server/Sessions/SessionManager.cs b/src/ZB.MOM.WW.MxGateway.Server/Sessions/SessionManager.cs
index 0ded52e..ed9172f 100644
--- a/src/ZB.MOM.WW.MxGateway.Server/Sessions/SessionManager.cs
+++ b/src/ZB.MOM.WW.MxGateway.Server/Sessions/SessionManager.cs
@@ -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);
+ }
+ }
+
///
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
diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/GatewaySessionTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/GatewaySessionTests.cs
index 8d07108..23c9aa1 100644
--- a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/GatewaySessionTests.cs
+++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/GatewaySessionTests.cs
@@ -76,6 +76,49 @@ public sealed class GatewaySessionTests
await session.DisposeAsync();
}
+ ///
+ /// A worker client that arrives after the session has already been closed (an admin kill
+ /// or the host drain landed while the worker was still starting) must be refused:
+ /// returns and
+ /// leaves unset, so the caller can discard the
+ /// client instead of handing a dead session back as an opened one.
+ ///
+ /// A task that represents the asynchronous operation.
+ [Fact]
+ public async Task AttachWorkerClient_AfterClose_RefusesLateWorkerClient()
+ {
+ GatewaySession session = CreateSession();
+ session.TransitionTo(SessionState.StartingWorker);
+
+ await session.CloseAsync("admin-kill", CancellationToken.None);
+ bool attached = session.AttachWorkerClient(new BlockingShutdownWorkerClient());
+
+ Assert.False(attached);
+ Assert.Null(session.WorkerClient);
+ Assert.Equal(SessionState.Closed, session.State);
+
+ await session.DisposeAsync();
+ }
+
+ ///
+ /// Disposal is the other end of the same window: once the session's distributor has been
+ /// torn down nothing can ever pump a late worker's events, so the attach is refused there
+ /// too rather than storing a client that is never read.
+ ///
+ /// A task that represents the asynchronous operation.
+ [Fact]
+ public async Task AttachWorkerClient_AfterDispose_RefusesLateWorkerClient()
+ {
+ GatewaySession session = CreateSession();
+ session.TransitionTo(SessionState.StartingWorker);
+
+ await session.DisposeAsync();
+ bool attached = session.AttachWorkerClient(new BlockingShutdownWorkerClient());
+
+ Assert.False(attached);
+ Assert.Null(session.WorkerClient);
+ }
+
///
/// A issued
/// while is parked between its
diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs
index 5528e9d..8e6c915 100644
--- a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs
+++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionManagerTests.cs
@@ -978,6 +978,99 @@ public sealed class SessionManagerTests
Assert.Equal(1, metrics.GetSnapshot().Faults);
}
+ ///
+ /// A failed open must leave no per-session event metrics behind: the rollback goes through
+ /// the same removal path every other teardown uses, so EventsBySession is cleared
+ /// for the session that never opened.
+ ///
+ /// A task that represents the asynchronous operation.
+ [Fact]
+ public async Task OpenSessionAsync_WhenWorkerCreationFails_RemovesSessionEventMetrics()
+ {
+ using GatewayMetrics metrics = new();
+ EventRecordingFailingSessionWorkerClientFactory factory = new(metrics);
+ SessionManager manager = CreateManager(factory, metrics: metrics);
+
+ await Assert.ThrowsAsync(
+ async () => await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None));
+
+ Assert.NotNull(factory.StartedSessionId);
+ Assert.False(metrics.GetSnapshot().EventsBySession.ContainsKey(factory.StartedSessionId!));
+ }
+
+ ///
+ /// The session is registered before its worker is created, so an admin Kill can land while
+ /// the worker is still starting (up to StartupTimeoutSeconds). The late worker client
+ /// must not be attached to the already-closed session and handed back to the caller as if
+ /// the open had succeeded: it is killed and disposed, the open fails with
+ /// OpenFailed, and nothing is left in the registry or the metrics.
+ ///
+ /// A task that represents the asynchronous operation.
+ [Fact]
+ public async Task OpenSessionAsync_WhenKilledDuringWorkerStartup_FailsOpenAndDiscardsLateWorkerClient()
+ {
+ FakeWorkerClient workerClient = new();
+ GatedSessionWorkerClientFactory factory = new(workerClient);
+ SessionRegistry registry = new();
+ using GatewayMetrics metrics = new();
+ SessionManager manager = CreateManager(factory, registry, metrics, CreateOptions(maxSessions: 1));
+
+ Task openTask = manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
+ GatewaySession starting = await factory.WaitForCreateStartAsync();
+
+ SessionCloseResult kill = await manager.KillWorkerAsync(starting.SessionId, "admin-kill", CancellationToken.None);
+ factory.ReleaseCreate();
+ SessionManagerException exception = await Assert.ThrowsAsync(async () => await openTask);
+
+ Assert.Equal(SessionState.Closed, kill.FinalState);
+ Assert.Equal(SessionManagerErrorCode.OpenFailed, exception.ErrorCode);
+ Assert.Equal(1, workerClient.KillCount);
+ Assert.Equal(1, workerClient.DisposeCount);
+ Assert.Null(starting.WorkerClient);
+ Assert.Equal(0, registry.Count);
+ GatewayMetricsSnapshot snapshot = metrics.GetSnapshot();
+ Assert.Equal(0, snapshot.OpenSessions);
+ Assert.Equal(0, snapshot.SessionsOpened);
+ Assert.Equal(0, snapshot.SessionsClosed);
+
+ // The slot the killed open held must have been released exactly once: a replacement open
+ // succeeds against MaxSessions = 1 (the released gate hands the same fake straight back).
+ GatewaySession reopened = await manager.OpenSessionAsync(CreateOpenRequest(), "client-2", ownerKeyId: null, CancellationToken.None);
+ Assert.Equal(SessionState.Ready, reopened.State);
+ Assert.Equal(1, registry.Count);
+ }
+
+ ///
+ /// Same window as the kill case, driven by the host drain: a session still starting its
+ /// worker when ShutdownAsync runs is closed by the drain, and the worker client that
+ /// arrives afterwards is discarded rather than left running as an orphan.
+ ///
+ /// A task that represents the asynchronous operation.
+ [Fact]
+ public async Task OpenSessionAsync_WhenShutdownDuringWorkerStartup_FailsOpenAndDiscardsLateWorkerClient()
+ {
+ FakeWorkerClient workerClient = new();
+ GatedSessionWorkerClientFactory factory = new(workerClient);
+ SessionRegistry registry = new();
+ using GatewayMetrics metrics = new();
+ SessionManager manager = CreateManager(factory, registry, metrics);
+
+ Task openTask = manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None);
+ GatewaySession starting = await factory.WaitForCreateStartAsync();
+
+ await manager.ShutdownAsync(CancellationToken.None);
+ factory.ReleaseCreate();
+ SessionManagerException exception = await Assert.ThrowsAsync(async () => await openTask);
+
+ Assert.Equal(SessionManagerErrorCode.OpenFailed, exception.ErrorCode);
+ Assert.Equal(SessionState.Closed, starting.State);
+ Assert.Equal(1, workerClient.KillCount);
+ Assert.Equal(1, workerClient.DisposeCount);
+ Assert.Equal(0, registry.Count);
+ Assert.Equal(0, metrics.GetSnapshot().OpenSessions);
+ Assert.Equal(0, metrics.GetSnapshot().SessionsClosed);
+ }
+
/// Verifies that closing expired leases only closes expired sessions.
/// A task that represents the asynchronous operation.
[Fact]
@@ -1422,6 +1515,61 @@ public sealed class SessionManagerTests
}
}
+ ///
+ /// Factory that records a per-session event metric for the session it is starting and
+ /// then fails, so a test can check the open rollback clears that metric.
+ ///
+ private sealed class EventRecordingFailingSessionWorkerClientFactory(GatewayMetrics metrics) : ISessionWorkerClientFactory
+ {
+ /// Gets the id of the session whose worker creation was attempted.
+ public string? StartedSessionId { get; private set; }
+
+ ///
+ public Task CreateAsync(
+ GatewaySession session,
+ CancellationToken cancellationToken)
+ {
+ StartedSessionId = session.SessionId;
+ metrics.EventReceived(session.SessionId, MxEventFamily.OnDataChange.ToString());
+ throw new InvalidOperationException("worker startup failed");
+ }
+ }
+
+ ///
+ /// Factory whose parks until the test releases it, so the test can
+ /// act on the registered-but-still-starting session (kill it, drain the manager) inside the
+ /// window between registry insertion and worker attach. Once released it hands the same
+ /// client to every later call.
+ ///
+ private sealed class GatedSessionWorkerClientFactory(IWorkerClient workerClient) : ISessionWorkerClientFactory
+ {
+ private readonly TaskCompletionSource _entered = new(TaskCreationOptions.RunContinuationsAsynchronously);
+ private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+ /// Waits until the manager has asked for a worker and returns the starting session.
+ /// The session whose worker is being created.
+ public Task WaitForCreateStartAsync()
+ {
+ return _entered.Task.WaitAsync(TimeSpan.FromSeconds(5));
+ }
+
+ /// Lets the parked worker creation complete.
+ public void ReleaseCreate()
+ {
+ _release.TrySetResult();
+ }
+
+ ///
+ public async Task CreateAsync(
+ GatewaySession session,
+ CancellationToken cancellationToken)
+ {
+ _entered.TrySetResult(session);
+ await _release.Task.WaitAsync(cancellationToken);
+ return workerClient;
+ }
+ }
+
private sealed class FakeWorkerClient : IWorkerClient
{
private WorkerClientState _state = WorkerClientState.Ready;