fix(worker): reply WorkerUnavailable to commands that race shutdown (#182)

TryStartCommandTask returned silently when command intake had stopped, so a
command refused during shutdown got no reply and no log and the gateway's
correlation wait ran out its own timeout. The refusal now writes a
WorkerUnavailable reply for the correlation and logs
WorkerCommandRejectedAfterShutdown. StopAcceptingCommands is internal so the
test can reach the guard (the serial read loop cannot race it today).

Test: WorkerPipeSessionTests.CommandAfterShutdown_RepliesWorkerUnavailable.
This commit is contained in:
Joseph Doherty
2026-11-06 16:03:47 -05:00
parent e5e7f1854f
commit 3a3d9e8e06
3 changed files with 131 additions and 5 deletions
+5 -1
View File
@@ -1011,7 +1011,11 @@ Graceful shutdown sequence:
1. Pipe reader receives `WorkerShutdown`. 1. Pipe reader receives `WorkerShutdown`.
2. Worker host marks shutdown requested. 2. Worker host marks shutdown requested.
3. Reject new commands. 3. Reject new commands. A command that reaches the dispatcher after intake has
stopped is answered with a `WorkerUnavailable` reply and logged as
`WorkerCommandRejectedAfterShutdown`, so the gateway's correlation wait ends
at once instead of running out its own timeout. The read loop is serial and
stops after the shutdown frame, so this guard is defensive today.
4. Let current STA command finish if within timeout. 4. Let current STA command finish if within timeout.
5. Optionally run MXAccess cleanup: 5. Optionally run MXAccess cleanup:
- `UnAdvise`, - `UnAdvise`,
@@ -1943,6 +1943,60 @@ public sealed class WorkerPipeSessionTests
&& (string?)correlationId == "command-dropped-after-shutdown"); && (string?)correlationId == "command-dropped-after-shutdown");
} }
/// <summary>
/// A command that reaches the dispatcher after shutdown has stopped command intake must be
/// answered with <see cref="ProtocolStatusCode.WorkerUnavailable"/> (and logged) rather than
/// dropped silently, which would leave the gateway's correlation wait to run out its own
/// timeout with no trace (#182).
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task CommandAfterShutdown_RepliesWorkerUnavailable()
{
using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
FakeRuntimeSession runtime = new();
RecordingWorkerLogger logger = new();
WorkerFrameProtocolOptions options = CreateOptions();
WorkerPipeSession session = new(
new WorkerFrameReader(pipePair.WorkerStream, options),
new WorkerFrameWriter(pipePair.WorkerStream, options),
options,
() => 1234,
new WorkerPipeSessionOptions
{
HeartbeatInterval = TimeSpan.FromSeconds(1),
HeartbeatGrace = TimeSpan.FromSeconds(5),
},
() => runtime,
logger);
Task runTask = session.RunAsync(cancellation.Token);
await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);
session.StopAcceptingCommands();
await pipePair.GatewayWriter.WriteAsync(
CreateCommandEnvelope("command-after-shutdown"),
cancellation.Token);
WorkerEnvelope reply = await ReadUntilAsync(
pipePair.GatewayReader,
WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
envelope => envelope.CorrelationId == "command-after-shutdown",
cancellation.Token);
Assert.Equal(ProtocolStatusCode.WorkerUnavailable, reply.WorkerCommandReply.Reply.ProtocolStatus.Code);
Assert.Equal(MxCommandKind.Register, reply.WorkerCommandReply.Reply.Kind);
Assert.Equal(SessionId, reply.WorkerCommandReply.Reply.SessionId);
Assert.False(runtime.DispatchStarted.IsSet);
Assert.Contains(
logger.Events,
entry => entry.EventName == "WorkerCommandRejectedAfterShutdown"
&& entry.Fields.TryGetValue("correlation_id", out object? correlationId)
&& (string?)correlationId == "command-after-shutdown");
await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
}
private static WorkerPipeSession CreateSession( private static WorkerPipeSession CreateSession(
Stream inbound, Stream inbound,
Stream outbound, Stream outbound,
@@ -748,7 +748,11 @@ public sealed class WorkerPipeSession
return await HandleControlCommandAsync(envelope, cancellationToken).ConfigureAwait(false); return await HandleControlCommandAsync(envelope, cancellationToken).ConfigureAwait(false);
} }
TryStartCommandTask(envelope, cancellationToken); if (!TryStartCommandTask(envelope, cancellationToken))
{
await RejectCommandAfterShutdownAsync(envelope, cancellationToken).ConfigureAwait(false);
}
return true; return true;
case WorkerEnvelope.BodyOneofCase.WorkerShutdown: case WorkerEnvelope.BodyOneofCase.WorkerShutdown:
await ShutdownAsync(envelope.WorkerShutdown, cancellationToken).ConfigureAwait(false); await ShutdownAsync(envelope.WorkerShutdown, cancellationToken).ConfigureAwait(false);
@@ -1230,7 +1234,12 @@ public sealed class WorkerPipeSession
} }
} }
private void TryStartCommandTask( /// <summary>
/// Starts the command on the runtime session unless shutdown has already stopped command
/// intake. Returns <c>false</c> when the command was refused, so the caller can answer the
/// correlation instead of leaving the gateway's wait to run out its own timeout.
/// </summary>
private bool TryStartCommandTask(
WorkerEnvelope envelope, WorkerEnvelope envelope,
CancellationToken cancellationToken) CancellationToken cancellationToken)
{ {
@@ -1239,7 +1248,7 @@ public sealed class WorkerPipeSession
{ {
if (!_acceptingCommands) if (!_acceptingCommands)
{ {
return; return false;
} }
commandTask = ProcessCommandAsync(envelope, cancellationToken); commandTask = ProcessCommandAsync(envelope, cancellationToken);
@@ -1254,6 +1263,60 @@ public sealed class WorkerPipeSession
// would fault the discarded task with no observer. Close the gap the same way the rest of // would fault the discarded task with no observer. Close the gap the same way the rest of
// the file does. // the file does.
ObserveFaultWhenever(ObserveCommandTaskAsync(commandTask)); ObserveFaultWhenever(ObserveCommandTaskAsync(commandTask));
return true;
}
/// <summary>
/// Answers a command that arrived after shutdown stopped command intake with a
/// <see cref="ProtocolStatusCode.WorkerUnavailable"/> reply, and logs it. The worker is
/// going away, so the reply is best effort: if the pipe is already unusable the failure is
/// logged and the gateway falls back to its own correlation timeout, as it did before.
/// </summary>
private async Task RejectCommandAfterShutdownAsync(
WorkerEnvelope envelope,
CancellationToken cancellationToken)
{
MxCommandKind kind = envelope.WorkerCommand?.Command?.Kind ?? MxCommandKind.Unspecified;
const string message = "The worker is shutting down and no longer accepts commands.";
_logger?.Information(
"WorkerCommandRejectedAfterShutdown",
new Dictionary<string, object?>
{
["correlation_id"] = envelope.CorrelationId,
["command_kind"] = kind.ToString(),
["worker_state"] = _state.ToString(),
});
MxCommandReply reply = new()
{
SessionId = _options.SessionId,
CorrelationId = envelope.CorrelationId,
Kind = kind,
DiagnosticMessage = message,
ProtocolStatus = new ProtocolStatus
{
Code = ProtocolStatusCode.WorkerUnavailable,
Message = message,
},
};
try
{
await WriteControlReplyAsync(reply, cancellationToken).ConfigureAwait(false);
}
catch (Exception exception) when (exception is IOException
|| exception is ObjectDisposedException
|| exception is WorkerFrameProtocolException)
{
_logger?.Error(
"WorkerCommandRejectionReplyFailed",
new Dictionary<string, object?>
{
["correlation_id"] = envelope.CorrelationId,
["command_kind"] = kind.ToString(),
["reason"] = exception.Message,
});
}
} }
private async Task ObserveCommandTaskAsync(Task commandTask) private async Task ObserveCommandTaskAsync(Task commandTask)
@@ -1274,7 +1337,12 @@ public sealed class WorkerPipeSession
} }
} }
private void StopAcceptingCommands() /// <summary>
/// Stops command intake. Internal so tests can reach the refusal path directly: the read
/// loop is serial and ends after a shutdown, so a command cannot race it over the pipe today,
/// but the guard must still answer if a future change lets one through.
/// </summary>
internal void StopAcceptingCommands()
{ {
lock (_commandTaskGate) lock (_commandTaskGate)
{ {