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:
mika kuns
2026-08-06 14:34:48 +02:00
parent bac8387069
commit 109e85da83
4 changed files with 180 additions and 15 deletions
@@ -47,6 +47,7 @@ public sealed class QueueServiceTests : IDisposable
private QueueWaker _waker = null!;
private FakeUsageGate _usageGate = null!;
private CapturingHubContext _hubContext = null!;
private RunCancellationRegistry _runCancels = null!;
private (QueueService service, FakeClaudeProcess fakeProcess) CreateService(
Func<string, string, IReadOnlyList<string>, Func<string, Task>, CancellationToken, Task<RunResult>>? handler = null,
@@ -67,6 +68,7 @@ public sealed class QueueServiceTests : IDisposable
var picker = new QueuePicker(dbFactory);
var overrideSlot = new OverrideSlotService(dbFactory, runner, NullLogger<OverrideSlotService>.Instance, built.RunCancels);
_usageGate = usageGate ?? new FakeUsageGate();
_runCancels = built.RunCancels;
var service = new QueueService(dbFactory, runner, _cfg, NullLogger<QueueService>.Instance, _waker, picker, overrideSlot, state, built.RunCancels,
_usageGate, usageState ?? new UsageState(), broadcaster);
return (service, fake);
@@ -619,4 +621,91 @@ public sealed class QueueServiceTests : IDisposable
Assert.False(cancelled);
cts.Cancel();
}
// Regression for the RunCancellationRegistry.Register-return double-dispatch bug: the
// queue picker claims Queued->Running atomically before it registers its CTS, so a
// concurrent RunNow that registered first can momentarily hold the registry slot for a
// task the picker is about to run. The picker must retry registration (letting the loser's
// own claim-attempt-then-unregister resolve) rather than dispatching under an unregistered
// CTS. Once it does register, TryCancel must be able to reach the actually-running process.
[Fact]
public async Task QueuePicker_RetriesRegistration_WhenRunNowLoserUnregistersInTime()
{
var listId = await SeedListAsync();
var task = await SeedQueuedTask(listId);
var running = new TaskCompletionSource();
var cancelled = false;
var (service, fake) = CreateService(async (_, _, _, _, ct) =>
{
running.SetResult();
try
{
await Task.Delay(Timeout.Infinite, ct);
}
catch (OperationCanceledException)
{
cancelled = true;
throw;
}
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
});
// Simulate a RunNow that registered first, then loses its own claim attempt against
// the row the picker is about to claim and cleans up shortly after (well within the
// picker's retry window).
var runNowCts = new CancellationTokenSource();
Assert.True(_runCancels.Register(task.Id, runNowCts));
_ = Task.Delay(60).ContinueWith(_ =>
{
_runCancels.Unregister(task.Id, runNowCts);
runNowCts.Dispose();
});
using var cts = new CancellationTokenSource();
await service.StartAsync(cts.Token);
_waker.Wake();
await running.Task.WaitAsync(TimeSpan.FromSeconds(5));
Assert.Equal(1, fake.CallCount);
// The picker's retry must have registered the CTS actually driving this run.
Assert.True(service.CancelTask(task.Id));
await Task.Delay(200);
Assert.True(cancelled);
cts.Cancel();
}
// If registration never resolves (the loser never unregisters), the picker must not run
// the task unregistered — it marks the already-claimed row Failed instead of leaving it
// stuck Running with no way to cancel it.
[Fact]
public async Task QueuePicker_RegistrationNeverResolves_MarksTaskFailed_WithoutRunning()
{
var listId = await SeedListAsync();
var task = await SeedQueuedTask(listId);
var (service, fake) = CreateService((_, _, _, _, _) =>
Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" }));
using var stuckCts = new CancellationTokenSource();
Assert.True(_runCancels.Register(task.Id, stuckCts));
using var cts = new CancellationTokenSource();
await service.StartAsync(cts.Token);
_waker.Wake();
// Max retry window is 10 * 20ms; give it comfortable headroom.
await Task.Delay(600);
cts.Cancel();
Assert.Equal(0, fake.CallCount);
var reloaded = await _taskRepo.GetByIdAsync(task.Id);
Assert.Equal(TaskStatus.Failed, reloaded!.Status);
// The still-registered "other" run must be untouched.
Assert.False(stuckCts.IsCancellationRequested);
}
}