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); + } }