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.
2026-10-22 06:16:53 -04:00
parent 488649ce99
commit 87f23d9712
5 changed files with 389 additions and 39 deletions
@@ -76,6 +76,49 @@ public sealed class GatewaySessionTests
        await session.DisposeAsync();
    }

+    /// <summary>
+    ///     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:
+    ///     <see cref="GatewaySession.AttachWorkerClient"/> returns <see langword="false"/> and
+    ///     leaves <see cref="GatewaySession.WorkerClient"/> unset, so the caller can discard the
+    ///     client instead of handing a dead session back as an opened one.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [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();
+    }
+
+    /// <summary>
+    ///     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.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [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);
+    }
+
    /// <summary>
    ///     A <see cref="GatewaySession.MarkFaulted"/> issued
    ///     while <see cref="GatewaySession.CloseAsync"/> is parked between its
@@ -978,6 +978,99 @@ public sealed class SessionManagerTests
        Assert.Equal(1, metrics.GetSnapshot().Faults);
    }

+    /// <summary>
+    ///     A failed open must leave no per-session event metrics behind: the rollback goes through
+    ///     the same removal path every other teardown uses, so <c>EventsBySession</c> is cleared
+    ///     for the session that never opened.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [Fact]
+    public async Task OpenSessionAsync_WhenWorkerCreationFails_RemovesSessionEventMetrics()
+    {
+        using GatewayMetrics metrics = new();
+        EventRecordingFailingSessionWorkerClientFactory factory = new(metrics);
+        SessionManager manager = CreateManager(factory, metrics: metrics);
+
+        await Assert.ThrowsAsync<SessionManagerException>(
+            async () => await manager.OpenSessionAsync(CreateOpenRequest(), "client-1", ownerKeyId: null, CancellationToken.None));
+
+        Assert.NotNull(factory.StartedSessionId);
+        Assert.False(metrics.GetSnapshot().EventsBySession.ContainsKey(factory.StartedSessionId!));
+    }
+
+    /// <summary>
+    ///     The session is registered before its worker is created, so an admin Kill can land while
+    ///     the worker is still starting (up to <c>StartupTimeoutSeconds</c>). 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
+    ///     <c>OpenFailed</c>, and nothing is left in the registry or the metrics.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [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<GatewaySession> 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<SessionManagerException>(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);
+    }
+
+    /// <summary>
+    ///     Same window as the kill case, driven by the host drain: a session still starting its
+    ///     worker when <c>ShutdownAsync</c> runs is closed by the drain, and the worker client that
+    ///     arrives afterwards is discarded rather than left running as an orphan.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [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<GatewaySession> 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<SessionManagerException>(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);
+    }
+
    /// <summary>Verifies that closing expired leases only closes expired sessions.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
@@ -1422,6 +1515,61 @@ public sealed class SessionManagerTests
        }
    }

+    /// <summary>
+    ///     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.
+    /// </summary>
+    private sealed class EventRecordingFailingSessionWorkerClientFactory(GatewayMetrics metrics) : ISessionWorkerClientFactory
+    {
+        /// <summary>Gets the id of the session whose worker creation was attempted.</summary>
+        public string? StartedSessionId { get; private set; }
+
+        /// <inheritdoc />
+        public Task<IWorkerClient> CreateAsync(
+            GatewaySession session,
+            CancellationToken cancellationToken)
+        {
+            StartedSessionId = session.SessionId;
+            metrics.EventReceived(session.SessionId, MxEventFamily.OnDataChange.ToString());
+            throw new InvalidOperationException("worker startup failed");
+        }
+    }
+
+    /// <summary>
+    ///     Factory whose <see cref="CreateAsync"/> 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.
+    /// </summary>
+    private sealed class GatedSessionWorkerClientFactory(IWorkerClient workerClient) : ISessionWorkerClientFactory
+    {
+        private readonly TaskCompletionSource<GatewaySession> _entered = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        /// <summary>Waits until the manager has asked for a worker and returns the starting session.</summary>
+        /// <returns>The session whose worker is being created.</returns>
+        public Task<GatewaySession> WaitForCreateStartAsync()
+        {
+            return _entered.Task.WaitAsync(TimeSpan.FromSeconds(5));
+        }
+
+        /// <summary>Lets the parked worker creation complete.</summary>
+        public void ReleaseCreate()
+        {
+            _release.TrySetResult();
+        }
+
+        /// <inheritdoc />
+        public async Task<IWorkerClient> 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;