From 109e85da83a498735b9f8ea85dfc3d524978fbc4 Mon Sep 17 00:00:00 2001 From: mika kuns Date: Thu, 6 Aug 2026 14:34:48 +0200 Subject: [PATCH] fix(worker): honor RunCancellationRegistry.Register's return value at both dispatch sites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../Queue/OverrideSlotService.cs | 8 +- src/ClaudeDo.Worker/Queue/QueueService.cs | 57 +++++++++--- .../Queue/OverrideSlotServiceTests.cs | 41 ++++++++- .../Services/QueueServiceTests.cs | 89 +++++++++++++++++++ 4 files changed, 180 insertions(+), 15 deletions(-) diff --git a/src/ClaudeDo.Worker/Queue/OverrideSlotService.cs b/src/ClaudeDo.Worker/Queue/OverrideSlotService.cs index 57454661..caac1419 100644 --- a/src/ClaudeDo.Worker/Queue/OverrideSlotService.cs +++ b/src/ClaudeDo.Worker/Queue/OverrideSlotService.cs @@ -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 => { diff --git a/src/ClaudeDo.Worker/Queue/QueueService.cs b/src/ClaudeDo.Worker/Queue/QueueService.cs index cdc6d427..b01a6ec9 100644 --- a/src/ClaudeDo.Worker/Queue/QueueService.cs +++ b/src/ClaudeDo.Worker/Queue/QueueService.cs @@ -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 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; diff --git a/tests/ClaudeDo.Worker.Tests/Queue/OverrideSlotServiceTests.cs b/tests/ClaudeDo.Worker.Tests/Queue/OverrideSlotServiceTests.cs index 5a437eaa..8e6f3f0a 100644 --- a/tests/ClaudeDo.Worker.Tests/Queue/OverrideSlotServiceTests.cs +++ b/tests/ClaudeDo.Worker.Tests/Queue/OverrideSlotServiceTests.cs @@ -31,7 +31,7 @@ public sealed class OverrideSlotServiceTests : IDisposable public void Dispose() { _db.Dispose(); try { Directory.Delete(_tempDir, true); } catch { } } - private OverrideSlotService BuildService() + private OverrideSlotService BuildService(RunCancellationRegistry? runCancels = null) { var dbFactory = _db.CreateFactory(); var state = TaskStateServiceBuilder.Build(dbFactory).State; @@ -39,7 +39,7 @@ public sealed class OverrideSlotServiceTests : IDisposable var runner = new TaskRunner(new FakeClaudeProcess(), dbFactory, new HubBroadcaster(new CapturingHubContext()), wt, new ClaudeArgsBuilder(), _cfg, NullLogger.Instance, state, new TaskRunTokenRegistry(), new AttachmentStore(), new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader()); - var runCancels = new RunCancellationRegistry(NullLogger.Instance); + runCancels ??= new RunCancellationRegistry(NullLogger.Instance); return new OverrideSlotService(dbFactory, runner, NullLogger.Instance, runCancels); } @@ -86,4 +86,41 @@ public sealed class OverrideSlotServiceTests : IDisposable // background, so this only asserts the precheck doesn't reject a claimable task. await service.RunNow(taskId); } + + // Regression for the RunCancellationRegistry.Register-return double-dispatch bug: the + // queue picker's atomic Queued->Running claim can land between RunNow's DB precheck and + // its registry registration. If the queue side already holds the registration, RunNow + // must fail loudly instead of silently registering over it (or being ignored and then + // unregistering the winner's CTS during its own cleanup). + [Fact] + public async Task RunNow_LosesRegistrationRaceToQueue_ThrowsAndLeavesWinnersCtsUntouched() + { + string listId = Guid.NewGuid().ToString(), taskId = Guid.NewGuid().ToString(); + using (var ctx = _db.CreateContext()) + { + ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow }); + ctx.Tasks.Add(new TaskEntity + { + Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Queued, + CreatedAt = DateTime.UtcNow, + }); + await ctx.SaveChangesAsync(); + } + + var runCancels = new RunCancellationRegistry(NullLogger.Instance); + var service = BuildService(runCancels); + + // Simulate the queue picker having already won the registration race for this task id. + using var winnerCts = new CancellationTokenSource(); + Assert.True(runCancels.Register(taskId, winnerCts)); + + var ex = await Assert.ThrowsAsync(() => service.RunNow(taskId)); + Assert.Contains("lost the double-dispatch race", ex.Message); + Assert.Null(service.CurrentSlot); + + // The loser must not have cancelled or displaced the winner's registration. + Assert.False(winnerCts.IsCancellationRequested); + Assert.True(runCancels.TryCancel(taskId)); + Assert.True(winnerCts.IsCancellationRequested); + } } diff --git a/tests/ClaudeDo.Worker.Tests/Services/QueueServiceTests.cs b/tests/ClaudeDo.Worker.Tests/Services/QueueServiceTests.cs index f9a28fae..e79d6819 100644 --- a/tests/ClaudeDo.Worker.Tests/Services/QueueServiceTests.cs +++ b/tests/ClaudeDo.Worker.Tests/Services/QueueServiceTests.cs @@ -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, Func, CancellationToken, Task>? handler = null, @@ -67,6 +68,7 @@ public sealed class QueueServiceTests : IDisposable var picker = new QueuePicker(dbFactory); var overrideSlot = new OverrideSlotService(dbFactory, runner, NullLogger.Instance, built.RunCancels); _usageGate = usageGate ?? new FakeUsageGate(); + _runCancels = built.RunCancels; var service = new QueueService(dbFactory, runner, _cfg, NullLogger.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); + } }