From b0100e738adbc055c6d11246abefc982d42acf4a Mon Sep 17 00:00:00 2001 From: Joseph Doherty Date: Thu, 22 Oct 2026 06:10:52 -0400 Subject: [PATCH] fix(worker): bound ShutdownAsync's envelope enqueue by the shutdown timeout WorkerClient.ShutdownAsync marked the client Closing and then queued the WorkerShutdown envelope on the caller's token alone; the shutdown-timeout CTS was only created afterwards. The outbound channel is bounded, and a worker that has stopped reading the pipe parks the write loop in a pipe write it can only leave through _stopCts while the channel behind it fills with commands and their cancels. Once Closing, the heartbeat watchdog skips every tick, so nothing cancelled _stopCts and the enqueue waited for the caller's token - the lease sweep's stopping token, which fires at host stop and never before - with GatewaySession.CloseAsync holding its close lock throughout. Create the linked timeout CTS before the enqueue and run the enqueue, the loop waits and the process-exit wait on it, so the whole call is bounded by the shutdown timeout. A stalled queue now surfaces as the existing ShutdownTimeout fault, which cancels the loops, kills the worker process tree and hands GatewaySession its kill fallback. Test: WorkerClientTests.ShutdownAsync_WhenOutboundQueueIsFullBehindStalledPipeWrite_TimesOutAndKillsWorker stalls the gateway-side stream (a worker that stopped reading), fills the outbound channel with abandoned commands and their cancels, and asserts that ShutdownAsync throws ShutdownTimeout inside the test bound and kills the fake worker. It hung past the bound before the fix. --- docs/GatewayProcessDesign.md | 14 +- docs/Sessions.md | 2 +- .../Workers/WorkerClient.cs | 17 +- .../Gateway/Workers/WorkerClientTests.cs | 203 +++++++++++++++++- 4 files changed, 228 insertions(+), 8 deletions(-) diff --git a/docs/GatewayProcessDesign.md b/docs/GatewayProcessDesign.md index 035a9e1..549c563 100644 --- a/docs/GatewayProcessDesign.md +++ b/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 -the timeout, the close operation fails with `ShutdownTimeout`; `GatewaySession` -then kills the worker process tree before surfacing the close failure. +exit within the same shutdown timeout. The timeout starts before the +`WorkerShutdown` envelope is queued, not after: the outbound channel is bounded, +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 diff --git a/docs/Sessions.md b/docs/Sessions.md index 7bcbc51..75a7f2d 100644 --- a/docs/Sessions.md +++ b/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) diff --git a/src/ZB.MOM.WW.MxGateway.Server/Workers/WorkerClient.cs b/src/ZB.MOM.WW.MxGateway.Server/Workers/WorkerClient.cs index 873dcdd..8cd9a63 100644 --- a/src/ZB.MOM.WW.MxGateway.Server/Workers/WorkerClient.cs +++ b/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"); diff --git a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerClientTests.cs b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerClientTests.cs index ec1824e..3585b30 100644 --- a/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerClientTests.cs +++ b/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerClientTests.cs @@ -707,6 +707,82 @@ public sealed class WorkerClientTests Assert.True(process.Disposed); } + /// + /// ShutdownAsync 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 _stopCts, and the bounded outbound channel + /// behind it fills with commands and their cancels. Once the client is Closing, the + /// heartbeat watchdog no longer fires (it only watches a Ready client) and nothing + /// else cancels _stopCts, 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 GatewaySession.CloseAsync holding its close lock the whole time. The shutdown + /// timeout has to cover the enqueue as well as the waits, so the close surfaces + /// and the worker is killed. + /// + /// A task that represents the asynchronous operation. + [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 invokeTask = client.InvokeAsync( + CreateCommand(MxCommandKind.Advise), + TestTimeout, + callerCts.Token); + if (round == 0) + { + await gatewayStream.FirstStalledWrite.WaitAsync(TestTimeout); + } + + await callerCts.CancelAsync(); + await Assert.ThrowsAnyAsync( + async () => await invokeTask.WaitAsync(TestTimeout)); + } + + TimeSpan shutdownTimeout = TimeSpan.FromMilliseconds(500); + DateTimeOffset startedAt = DateTimeOffset.UtcNow; + WorkerClientException exception = await Assert.ThrowsAsync( + 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); + } + /// /// Disposing a still-Ready client while the worker is emitting events must be /// an ordinary teardown, not a fault. DisposeAsync 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 } } + /// + /// 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. + /// + /// + /// 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 , + /// so that is the only write overload that stalls. + /// + 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; + + /// Completes when the first write parks after . + 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(); + } + + /// Makes every subsequent write park until cancelled or disposed. + public void StallWrites() => _stalled = true; + + public override async ValueTask WriteAsync( + ReadOnlyMemory 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(buffer, offset, count), cancellationToken).AsTask(); + } + + public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + + public override ValueTask ReadAsync( + Memory buffer, + CancellationToken cancellationToken = default) + { + return inner.ReadAsync(buffer, cancellationToken); + } + + public override Task 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); + } + } }