fix(worker): honor RunCancellationRegistry.Register's return value at both dispatch sites
Register(taskId, cts) already refuses (and logs) a double-registration, but both call sites discarded the bool and dispatched anyway under an unregistered CTS. If the loser then unregistered the winner's CTS during its own cleanup, TryCancel could silently no-op against a live process. - OverrideSlotService.StartInSlot now fails RunNow/ContinueTask loudly (throws) when it loses the registration race instead of registering over — or silently proceeding despite losing to — the queue picker. - QueueService's picker loop retries registration briefly (the loser's own claim-attempt-then-unregister resolves fast) before dispatching; if registration never resolves it marks the already-claimed row Failed instead of running it unregistered. - RunCancellationRegistry.Unregister already had compare-and-remove semantics (TryRemove(KeyValuePair)), so a loser's cleanup could not have removed the winner's CTS once registration correctly failed. Added regression tests exercising the real registry through both dispatch paths: RunNow losing the registration race throws without disturbing the winner, the picker's retry succeeds and TryCancel reaches the live run when the loser unregisters in time, and the picker fails the task instead of running unregistered when it never does.
This commit is contained in:
@@ -74,8 +74,14 @@ public sealed class OverrideSlotService
|
||||
throw new InvalidOperationException("override slot busy");
|
||||
|
||||
var cts = new CancellationTokenSource();
|
||||
if (!_runCancels.Register(taskId, cts))
|
||||
{
|
||||
cts.Dispose();
|
||||
throw new InvalidOperationException(
|
||||
$"Task '{taskId}' lost the double-dispatch race to the queue picker; it is already running there.");
|
||||
}
|
||||
|
||||
_slot = new QueueSlotState { TaskId = taskId, StartedAt = DateTime.UtcNow, Cts = cts };
|
||||
_runCancels.Register(taskId, cts);
|
||||
|
||||
_ = work(cts.Token).ContinueWith(t =>
|
||||
{
|
||||
|
||||
@@ -145,22 +145,38 @@ public sealed class QueueService : BackgroundService
|
||||
var task = await _picker.ClaimNextAsync(DateTime.UtcNow, stoppingToken);
|
||||
if (task is null) break;
|
||||
|
||||
// The row is already claimed (Queued->Running) here, but a concurrent
|
||||
// RunNow for the same task id may have registered its CTS first. Retry
|
||||
// briefly rather than dispatching under an unregistered CTS — the loser
|
||||
// of that race (TaskRunner.RunAsync's own claim attempt fails against the
|
||||
// row we just claimed) unregisters quickly once it bails.
|
||||
var cts = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken);
|
||||
if (!await RegisterWithRetryAsync(task.Id, cts, stoppingToken))
|
||||
{
|
||||
cts.Dispose();
|
||||
_logger.LogError(
|
||||
"Task {TaskId} claimed by the queue picker but could not be registered in " +
|
||||
"RunCancellationRegistry (double-dispatch race with RunNow did not resolve); " +
|
||||
"marking it Failed instead of running it unregistered.", task.Id);
|
||||
await _state.FailAsync(task.Id, DateTime.UtcNow,
|
||||
"Internal error: run cancellation registration contention.", CancellationToken.None);
|
||||
continue;
|
||||
}
|
||||
|
||||
lock (_lock)
|
||||
{
|
||||
var cts = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken);
|
||||
_queueSlots[task.Id] = new QueueSlotState { TaskId = task.Id, StartedAt = DateTime.UtcNow, Cts = cts };
|
||||
_runCancels.Register(task.Id, cts);
|
||||
|
||||
_ = RunInSlotAsync(task.Id, cts.Token).ContinueWith(t =>
|
||||
{
|
||||
if (t.IsFaulted)
|
||||
_logger.LogError(t.Exception, "RunInSlotAsync failed for task {TaskId} in queue slot", task.Id);
|
||||
lock (_lock) { _queueSlots.Remove(task.Id); }
|
||||
_runCancels.Unregister(task.Id, cts);
|
||||
cts.Dispose();
|
||||
_waker.Wake(); // Check for next task immediately.
|
||||
}, TaskScheduler.Default);
|
||||
}
|
||||
|
||||
_ = RunInSlotAsync(task.Id, cts.Token).ContinueWith(t =>
|
||||
{
|
||||
if (t.IsFaulted)
|
||||
_logger.LogError(t.Exception, "RunInSlotAsync failed for task {TaskId} in queue slot", task.Id);
|
||||
lock (_lock) { _queueSlots.Remove(task.Id); }
|
||||
_runCancels.Unregister(task.Id, cts);
|
||||
cts.Dispose();
|
||||
_waker.Wake(); // Check for next task immediately.
|
||||
}, TaskScheduler.Default);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -177,6 +193,23 @@ public sealed class QueueService : BackgroundService
|
||||
_logger.LogInformation("QueueService stopping");
|
||||
}
|
||||
|
||||
// A losing RunNow registers its CTS before the picker's atomic claim lands, then bails
|
||||
// (and unregisters) fast once its own claim attempt sees the row already Running. A few
|
||||
// short retries cover that window without stalling the picker loop indefinitely.
|
||||
private async Task<bool> RegisterWithRetryAsync(string taskId, CancellationTokenSource cts, CancellationToken ct)
|
||||
{
|
||||
const int maxAttempts = 10;
|
||||
const int delayMs = 20;
|
||||
|
||||
for (var attempt = 1; attempt <= maxAttempts; attempt++)
|
||||
{
|
||||
if (_runCancels.Register(taskId, cts)) return true;
|
||||
if (attempt < maxAttempts)
|
||||
await Task.Delay(delayMs, ct);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
private async Task ReportUsageGateTransitionAsync(UsageGateDecision decision)
|
||||
{
|
||||
if (decision.IsBlocked == _usageGateBlocked) return;
|
||||
|
||||
Reference in New Issue
Block a user