Merge stream B (bound WorkerClient.ShutdownAsync enqueue by shutdown timeout)

docs/GatewayProcessDesign.md

@@ -177,9 +177,17 @@ return a successful close result with the final known state.
 
 `WorkerClient.ShutdownAsync` sends `WorkerShutdown`, waits for the worker read,
 write, and heartbeat loops to stop, and waits for the launched worker process to
-exit within the same shutdown timeout. If the pipe loops or process exit exceed
+exit within the same shutdown timeout. The timeout starts before the
-the timeout, the close operation fails with `ShutdownTimeout`; `GatewaySession`
+`WorkerShutdown` envelope is queued, not after: the outbound channel is bounded,
-then kills the worker process tree before surfacing the close failure.
+and a worker that has stopped reading the pipe parks the write loop in a pipe
+write while the channel behind it fills with commands and their cancels. Because
+the heartbeat watchdog only watches a `Ready` client, nothing else would end that
+stall once the client is `Closing`, so an enqueue bound to the caller's token
+alone waited for that token — for the lease sweep, until host stop — with the
+session's close lock held. If queuing the envelope, the pipe loops, or process
+exit exceed the timeout, the close operation fails with `ShutdownTimeout` (the
+fault cancels the loops and kills the worker); `GatewaySession` then runs its
+kill fallback before surfacing the close failure.
 
 ### Invoke
 

docs/Sessions.md

@@ -267,7 +267,7 @@ The ReplayGap sentinel is emitted **only** on the `StreamEvents` server stream a
 
 ### Close
 
-`GatewaySession.CloseAsync` is serialized by a per-session `SemaphoreSlim` (`_closeLock`) so only one close runs at a time, but every read/write of `_state` still passes through `_syncRoot` (via `TryBeginClose` and `MarkClosed`). The close path therefore obeys the same lock discipline as `TransitionTo` / `MarkFaulted`: it transitions to `Closing`, asks the worker client to shut down within `ShutdownTimeout`, and on success transitions to `Closed`. `DisposeAsync` waits on `_closeLock` once before disposing the semaphore so an in-flight close's `Release()` cannot race against the dispose. If `WorkerClient.ShutdownAsync` throws, the session falls back to `IWorkerClient.Kill` (forced close):
+`GatewaySession.CloseAsync` is serialized by a per-session `SemaphoreSlim` (`_closeLock`) so only one close runs at a time, but every read/write of `_state` still passes through `_syncRoot` (via `TryBeginClose` and `MarkClosed`). The close path therefore obeys the same lock discipline as `TransitionTo` / `MarkFaulted`: it transitions to `Closing`, asks the worker client to shut down within `ShutdownTimeout`, and on success transitions to `Closed`. `DisposeAsync` waits on `_closeLock` once before disposing the semaphore so an in-flight close's `Release()` cannot race against the dispose. `WorkerClient.ShutdownAsync` is bounded by `ShutdownTimeout` end to end — queuing the `WorkerShutdown` envelope included, since a worker that stopped reading the pipe can leave the bounded outbound channel full with no watchdog left to clear it once the client is `Closing` — so a close never holds `_closeLock` longer than the timeout before reaching the fallback. If `WorkerClient.ShutdownAsync` throws, the session falls back to `IWorkerClient.Kill` (forced close):
 
 ```csharp
 if (_workerClient is not null)

src/ZB.MOM.WW.MxGateway.Server/Workers/WorkerClient.cs

@@ -382,13 +382,26 @@ public sealed class WorkerClient : IWorkerClient
         }
 
         MarkClosing();
-        await EnqueueAsync(CreateShutdownEnvelope(timeout, "gateway-shutdown"), cancellationToken).ConfigureAwait(false);
-        _outboundEnvelopes.Writer.TryComplete();
 
+        // The timeout governs the enqueue as well as the waits. Queuing the shutdown envelope is
+        // not instantaneous: _outboundEnvelopes is bounded, and a worker that has stopped reading
+        // the pipe leaves the write loop parked in a pipe write it can only leave through _stopCts,
+        // with the channel behind it filling up from commands and their cancels. Once the client is
+        // Closing the heartbeat watchdog no longer fires (it watches a Ready client only), so nothing
+        // else cancels _stopCts; an enqueue bound to the caller's token alone therefore waited for
+        // that token — the lease sweep's stopping token, which fires at host stop and never before —
+        // while GatewaySession.CloseAsync held its close lock. Running the enqueue on the same
+        // timeout token lets a stalled queue surface as the ShutdownTimeout below, whose fault
+        // cancels _stopCts, kills the worker and hands GatewaySession its kill fallback, instead of
+        // wedging the close. Completing the writer after the enqueue is unchanged: a shutdown that
+        // never made it into the queue is completed on the fault path.
         using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
         timeoutCts.CancelAfter(timeout);
         try
         {
+            await EnqueueAsync(CreateShutdownEnvelope(timeout, "gateway-shutdown"), timeoutCts.Token).ConfigureAwait(false);
+            _outboundEnvelopes.Writer.TryComplete();
+
             await WaitForBackgroundTasksAsync(timeoutCts.Token).ConfigureAwait(false);
             await WaitForProcessExitAsync(timeoutCts.Token).ConfigureAwait(false);
             MarkClosed("shutdown");

src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerClientTests.cs

@@ -707,6 +707,82 @@ public sealed class WorkerClientTests
         Assert.True(process.Disposed);
     }
 
+    /// <summary>
+    ///     <c>ShutdownAsync</c> must come back within its timeout even when the graceful request
+    ///     cannot be queued. A worker that stops reading the pipe leaves the write loop parked in
+    ///     a pipe write it can only leave through <c>_stopCts</c>, and the bounded outbound channel
+    ///     behind it fills with commands and their cancels. Once the client is <c>Closing</c>, the
+    ///     heartbeat watchdog no longer fires (it only watches a <c>Ready</c> client) and nothing
+    ///     else cancels <c>_stopCts</c>, so queuing the shutdown envelope on the caller's token
+    ///     alone blocked until that token fired — for the lease sweep, never before host stop —
+    ///     with <c>GatewaySession.CloseAsync</c> holding its close lock the whole time. The shutdown
+    ///     timeout has to cover the enqueue as well as the waits, so the close surfaces
+    ///     <see cref="WorkerClientErrorCode.ShutdownTimeout"/> and the worker is killed.
+    /// </summary>
+    /// <returns>A task that represents the asynchronous operation.</returns>
+    [Fact]
+    public async Task ShutdownAsync_WhenOutboundQueueIsFullBehindStalledPipeWrite_TimesOutAndKillsWorker()
+    {
+        await using PipePair pipePair = await PipePair.CreateAsync();
+        await using StallableWriteStream gatewayStream = new(pipePair.GatewayStream);
+        FakeWorkerProcess process = new(WorkerProcessId);
+
+        // One slot keeps the outbound bound small (MaxPendingCommands + 4 = 5 envelopes) so a
+        // handful of abandoned commands is enough to fill it behind the stalled write.
+        WorkerClientOptions options = new()
+        {
+            MaxPendingCommands = 1,
+            HeartbeatGrace = TimeSpan.FromSeconds(30),
+            HeartbeatCheckInterval = TimeSpan.FromSeconds(30),
+        };
+        await using WorkerClient client = CreateClient(
+            pipePair,
+            options,
+            processHandle: CreateProcessHandle(process),
+            gatewayStream: gatewayStream);
+        await CompleteHandshakeAsync(client, pipePair);
+
+        // The worker stops draining the pipe after the handshake: the next envelope the write loop
+        // dequeues parks it in WriteAsync until the stop token fires.
+        gatewayStream.StallWrites();
+
+        // Each abandoned command leaves two envelopes behind — the command the write loop already
+        // dequeued (or is still queued) and the WorkerCancel its abandonment forwards. The first
+        // command is what the write loop is now stuck on; the remaining five envelopes fill the
+        // channel to its bound.
+        for (int round = 0; round < 3; round++)
+        {
+            using CancellationTokenSource callerCts = new();
+            Task<WorkerCommandReply> invokeTask = client.InvokeAsync(
+                CreateCommand(MxCommandKind.Advise),
+                TestTimeout,
+                callerCts.Token);
+            if (round == 0)
+            {
+                await gatewayStream.FirstStalledWrite.WaitAsync(TestTimeout);
+            }
+
+            await callerCts.CancelAsync();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(
+                async () => await invokeTask.WaitAsync(TestTimeout));
+        }
+
+        TimeSpan shutdownTimeout = TimeSpan.FromMilliseconds(500);
+        DateTimeOffset startedAt = DateTimeOffset.UtcNow;
+        WorkerClientException exception = await Assert.ThrowsAsync<WorkerClientException>(
+            async () => await client.ShutdownAsync(shutdownTimeout, CancellationToken.None).WaitAsync(TestTimeout));
+        TimeSpan elapsed = DateTimeOffset.UtcNow - startedAt;
+
+        Assert.Equal(WorkerClientErrorCode.ShutdownTimeout, exception.ErrorCode);
+        Assert.True(
+            elapsed < TestTimeout,
+            $"ShutdownAsync took {elapsed.TotalMilliseconds:N0}ms against a {shutdownTimeout.TotalMilliseconds:N0}ms timeout.");
+        Assert.Equal(WorkerClientState.Faulted, client.State);
+        Assert.Equal(1, process.KillCount);
+        Assert.True(process.KillEntireProcessTree);
+        Assert.True(process.HasExited);
+    }
+
     /// <summary>
     ///     Disposing a still-<c>Ready</c> client while the worker is emitting events must be
     ///     an ordinary teardown, not a fault. <c>DisposeAsync</c> completes the same three channels
@@ -1186,7 +1262,8 @@ public sealed class WorkerClientTests
         GatewayMetrics? metrics = null,
         WorkerProcessHandle? processHandle = null,
         TimeProvider? timeProvider = null,
-        int maxMessageBytes = WorkerFrameProtocolOptions.DefaultMaxMessageBytes)
+        int maxMessageBytes = WorkerFrameProtocolOptions.DefaultMaxMessageBytes,
+        Stream? gatewayStream = null)
     {
         WorkerFrameProtocolOptions frameOptions = new(
             SessionId,
@@ -1195,7 +1272,7 @@ public sealed class WorkerClientTests
         WorkerClientConnection connection = new(
             SessionId,
             Nonce,
-            pipePair.GatewayStream,
+            gatewayStream ?? pipePair.GatewayStream,
             frameOptions,
             processHandle);
 
@@ -1478,4 +1555,126 @@ public sealed class WorkerClientTests
         }
     }
 
+    /// <summary>
+    ///     Gateway-side stream whose writes can be made to park exactly as a named-pipe write does
+    ///     once the worker has stopped reading and the OS buffer is full: the write completes only
+    ///     when its cancellation token fires or the stream is disposed. Reads pass through to the
+    ///     real pipe so the handshake and worker-side frames are unaffected.
+    /// </summary>
+    /// <remarks>
+    ///     Stalling in the stream rather than by overfilling the pipe keeps the test independent of
+    ///     the per-platform pipe buffer (64 KiB declared on Windows, a few KiB of socket buffer on
+    ///     macOS, ~200 KiB on Linux) and of how many frames it takes to exhaust it. The one call the
+    ///     frame writer makes is <see cref="WriteAsync(ReadOnlyMemory{byte}, CancellationToken)"/>,
+    ///     so that is the only write overload that stalls.
+    /// </remarks>
+    private sealed class StallableWriteStream(Stream inner) : Stream
+    {
+        private readonly CancellationTokenSource _disposed = new();
+        private readonly TaskCompletionSource _firstStalledWrite = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private volatile bool _stalled;
+        private bool _isDisposed;
+
+        /// <summary>Completes when the first write parks after <see cref="StallWrites"/>.</summary>
+        public Task FirstStalledWrite => _firstStalledWrite.Task;
+
+        public override bool CanRead => inner.CanRead;
+
+        public override bool CanSeek => false;
+
+        public override bool CanWrite => inner.CanWrite;
+
+        public override long Length => throw new NotSupportedException();
+
+        public override long Position
+        {
+            get => throw new NotSupportedException();
+            set => throw new NotSupportedException();
+        }
+
+        /// <summary>Makes every subsequent write park until cancelled or disposed.</summary>
+        public void StallWrites() => _stalled = true;
+
+        public override async ValueTask WriteAsync(
+            ReadOnlyMemory<byte> buffer,
+            CancellationToken cancellationToken = default)
+        {
+            if (!_stalled)
+            {
+                await inner.WriteAsync(buffer, cancellationToken);
+                return;
+            }
+
+            _firstStalledWrite.TrySetResult();
+            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(
+                cancellationToken,
+                _disposed.Token);
+            await Task.Delay(Timeout.Infinite, linked.Token);
+        }
+
+        public override Task WriteAsync(
+            byte[] buffer,
+            int offset,
+            int count,
+            CancellationToken cancellationToken)
+        {
+            return WriteAsync(new ReadOnlyMemory<byte>(buffer, offset, count), cancellationToken).AsTask();
+        }
+
+        public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+
+        public override ValueTask<int> ReadAsync(
+            Memory<byte> buffer,
+            CancellationToken cancellationToken = default)
+        {
+            return inner.ReadAsync(buffer, cancellationToken);
+        }
+
+        public override Task<int> ReadAsync(
+            byte[] buffer,
+            int offset,
+            int count,
+            CancellationToken cancellationToken)
+        {
+            return inner.ReadAsync(buffer, offset, count, cancellationToken);
+        }
+
+        public override int Read(byte[] buffer, int offset, int count) => inner.Read(buffer, offset, count);
+
+        public override void Flush() => inner.Flush();
+
+        public override Task FlushAsync(CancellationToken cancellationToken) => inner.FlushAsync(cancellationToken);
+
+        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+
+        public override void SetLength(long value) => throw new NotSupportedException();
+
+        // Disposed by the worker client (it owns the connection stream), then again by the test's
+        // await-using, so the release must be idempotent.
+        public override async ValueTask DisposeAsync()
+        {
+            if (!_isDisposed)
+            {
+                _isDisposed = true;
+                await _disposed.CancelAsync();
+                await inner.DisposeAsync();
+                _disposed.Dispose();
+            }
+
+            await base.DisposeAsync();
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (disposing && !_isDisposed)
+            {
+                _isDisposed = true;
+                _disposed.Cancel();
+                inner.Dispose();
+                _disposed.Dispose();
+            }
+
+            base.Dispose(disposing);
+        }
+    }
 }