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

This commit is contained in:
Joseph Doherty
2026-10-22 06:11:17 -04:00
4 changed files with 228 additions and 8 deletions
+11 -3
View File
@@ -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
+1 -1
View File
@@ -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)
@@ -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");
@@ -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);
}
}
}