mika kuns
109e85da83
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.
712 lines
25 KiB
C#
712 lines
25 KiB
C#
using ClaudeDo.Data;
|
|
using ClaudeDo.Data.Git;
|
|
using ClaudeDo.Data.Models;
|
|
using ClaudeDo.Data.Repositories;
|
|
using ClaudeDo.Worker.Config;
|
|
using ClaudeDo.Worker.Hub;
|
|
using ClaudeDo.Worker.Queue;
|
|
using ClaudeDo.Worker.Runner;
|
|
using ClaudeDo.Worker.Tests.Infrastructure;
|
|
using ClaudeDo.Worker.Usage;
|
|
using Microsoft.Extensions.Logging.Abstractions;
|
|
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;
|
|
|
|
namespace ClaudeDo.Worker.Tests.Services;
|
|
|
|
public sealed class QueueServiceTests : IDisposable
|
|
{
|
|
private readonly DbFixture _db = new();
|
|
private readonly ClaudeDoDbContext _ctx;
|
|
private readonly TaskRepository _taskRepo;
|
|
private readonly ListRepository _listRepo;
|
|
private readonly WorkerConfig _cfg;
|
|
private readonly string _tempDir;
|
|
|
|
public QueueServiceTests()
|
|
{
|
|
_ctx = _db.CreateContext();
|
|
_taskRepo = new TaskRepository(_ctx);
|
|
_listRepo = new ListRepository(_ctx);
|
|
_tempDir = Path.Combine(Path.GetTempPath(), $"claudedo_test_{Guid.NewGuid():N}");
|
|
Directory.CreateDirectory(_tempDir);
|
|
_cfg = new WorkerConfig
|
|
{
|
|
SandboxRoot = Path.Combine(_tempDir, "sandbox"),
|
|
LogRoot = Path.Combine(_tempDir, "logs"),
|
|
QueueBackstopIntervalMs = 50, // fast for tests
|
|
};
|
|
}
|
|
|
|
public void Dispose()
|
|
{
|
|
_ctx.Dispose();
|
|
_db.Dispose();
|
|
try { Directory.Delete(_tempDir, true); } catch { }
|
|
}
|
|
|
|
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,
|
|
FakeUsageGate? usageGate = null,
|
|
UsageState? usageState = null)
|
|
{
|
|
var fake = new FakeClaudeProcess(handler);
|
|
_hubContext = new CapturingHubContext();
|
|
var broadcaster = new HubBroadcaster(_hubContext);
|
|
var dbFactory = _db.CreateFactory();
|
|
var wtManager = new WorktreeManager(new GitService(), dbFactory, _cfg, NullLogger<WorktreeManager>.Instance);
|
|
var argsBuilder = new ClaudeArgsBuilder();
|
|
var built = TaskStateServiceBuilder.Build(dbFactory);
|
|
var state = built.State;
|
|
var runner = new TaskRunner(fake, dbFactory, broadcaster, wtManager, argsBuilder, _cfg,
|
|
NullLogger<TaskRunner>.Instance, state, new TaskRunTokenRegistry(), new AttachmentStore(), new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader());
|
|
_waker = new QueueWaker();
|
|
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);
|
|
}
|
|
|
|
private async Task SetAppSettingsAsync(
|
|
int maxParallel, int softPct = 50, int hardPct = 65, int gateFive = 80, int gateSeven = 90)
|
|
{
|
|
using var ctx = _db.CreateContext();
|
|
var repo = new AppSettingsRepository(ctx);
|
|
var settings = await repo.GetAsync();
|
|
settings.MaxParallelExecutions = maxParallel;
|
|
settings.UsageThrottleSoftPct = softPct;
|
|
settings.UsageThrottleHardPct = hardPct;
|
|
settings.UsageGateFiveHourPct = gateFive;
|
|
settings.UsageGateSevenDayPct = gateSeven;
|
|
await repo.UpdateAsync(settings);
|
|
}
|
|
|
|
private async Task<string> SeedListAsync()
|
|
{
|
|
var listId = Guid.NewGuid().ToString();
|
|
await _listRepo.AddAsync(new ListEntity { Id = listId, Name = "Test", CreatedAt = DateTime.UtcNow });
|
|
return listId;
|
|
}
|
|
|
|
private async Task<TaskEntity> SeedQueuedTask(string listId, DateTime? scheduledFor = null, DateTime? createdAt = null)
|
|
{
|
|
var task = new TaskEntity
|
|
{
|
|
Id = Guid.NewGuid().ToString(),
|
|
ListId = listId,
|
|
Title = "Test task",
|
|
Description = "Do something",
|
|
Status = TaskStatus.Queued,
|
|
ScheduledFor = scheduledFor,
|
|
CreatedAt = createdAt ?? DateTime.UtcNow,
|
|
};
|
|
await _taskRepo.AddAsync(task);
|
|
return task;
|
|
}
|
|
|
|
[Fact]
|
|
public async Task RunNow_Throws_When_Override_Slot_Busy()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
var tcs = new TaskCompletionSource<RunResult>();
|
|
|
|
var (service, _) = CreateService((_, _, _, _, ct) => tcs.Task);
|
|
|
|
var task1 = await SeedQueuedTask(listId);
|
|
var task2 = await SeedQueuedTask(listId);
|
|
|
|
await service.RunNow(task1.Id);
|
|
|
|
var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => service.RunNow(task2.Id));
|
|
Assert.Equal("override slot busy", ex.Message);
|
|
|
|
tcs.SetResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" });
|
|
}
|
|
|
|
[Fact]
|
|
public async Task RunNow_Throws_For_Unknown_Task()
|
|
{
|
|
var (service, _) = CreateService();
|
|
await Assert.ThrowsAsync<KeyNotFoundException>(() => service.RunNow("nonexistent"));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ReQueuedReviewTask_ResumesSession_WithFeedbackPrompt_AndClearsFeedback()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
|
|
IReadOnlyList<string>? capturedArgs = null;
|
|
string? capturedPrompt = null;
|
|
var done = new TaskCompletionSource();
|
|
|
|
var (service, _) = CreateService((prompt, _, args, _, _) =>
|
|
{
|
|
capturedPrompt = prompt;
|
|
capturedArgs = args;
|
|
done.TrySetResult();
|
|
return Task.FromResult(new RunResult { ExitCode = 0, SessionId = "sess-2", ResultMarkdown = "ok" });
|
|
});
|
|
|
|
// A task that was reviewed and rejected: Queued + ReviewFeedback, with a prior run carrying a session id.
|
|
var task = new TaskEntity
|
|
{
|
|
Id = Guid.NewGuid().ToString(),
|
|
ListId = listId,
|
|
Title = "Reviewed task",
|
|
Status = TaskStatus.Queued,
|
|
ReviewFeedback = "fix the bug",
|
|
CreatedAt = DateTime.UtcNow,
|
|
};
|
|
await _taskRepo.AddAsync(task);
|
|
await new TaskRunRepository(_ctx).AddAsync(new TaskRunEntity
|
|
{
|
|
Id = Guid.NewGuid().ToString(),
|
|
TaskId = task.Id,
|
|
RunNumber = 1,
|
|
IsRetry = false,
|
|
Prompt = "original",
|
|
SessionId = "sess-1",
|
|
StartedAt = DateTime.UtcNow.AddMinutes(-1),
|
|
});
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
await done.Task.WaitAsync(TimeSpan.FromSeconds(5));
|
|
|
|
Assert.NotNull(capturedArgs);
|
|
Assert.Contains("--resume", capturedArgs);
|
|
Assert.Contains("sess-1", capturedArgs);
|
|
Assert.Equal("fix the bug", capturedPrompt);
|
|
|
|
// Feedback is cleared after the run reaches a successful terminal state (post-run),
|
|
// so poll rather than asserting on the handler-fired instant.
|
|
var deadline = DateTime.UtcNow.AddSeconds(5);
|
|
TaskEntity? reloaded;
|
|
do
|
|
{
|
|
reloaded = await new TaskRepository(_db.CreateContext()).GetByIdAsync(task.Id);
|
|
if (reloaded?.ReviewFeedback is null) break;
|
|
await Task.Delay(25);
|
|
} while (DateTime.UtcNow < deadline);
|
|
cts.Cancel();
|
|
|
|
Assert.Equal(TaskStatus.WaitingForReview, reloaded!.Status);
|
|
Assert.Null(reloaded.ReviewFeedback);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Schedule_Filter_Skips_Future_Tasks()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId, scheduledFor: DateTime.UtcNow.AddHours(1));
|
|
|
|
var (service, fake) = CreateService((_, _, _, _, _) =>
|
|
Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" }));
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
|
|
// Start the service loop, wake it, give it time.
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
await Task.Delay(200);
|
|
cts.Cancel();
|
|
|
|
// The fake should never have been called because the task is scheduled in the future.
|
|
Assert.Equal(0, fake.CallCount);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Queue_FIFO_Sequentiality()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
|
|
var order = new List<string>();
|
|
var gate1 = new TaskCompletionSource();
|
|
var gate2 = new TaskCompletionSource();
|
|
var callCount = 0;
|
|
|
|
var (service, _) = CreateService(async (_, _, _, _, ct) =>
|
|
{
|
|
var n = Interlocked.Increment(ref callCount);
|
|
lock (order) { order.Add(n.ToString()); }
|
|
if (n == 1) await gate1.Task;
|
|
if (n == 2) gate2.SetResult();
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
});
|
|
|
|
await SeedQueuedTask(listId, createdAt: DateTime.UtcNow.AddSeconds(-2));
|
|
await SeedQueuedTask(listId, createdAt: DateTime.UtcNow.AddSeconds(-1));
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
|
|
// Wait until task1 has been picked up (poll instead of fixed delay to avoid flake under load).
|
|
var deadline = DateTime.UtcNow.AddSeconds(5);
|
|
while (order.Count == 0 && DateTime.UtcNow < deadline)
|
|
await Task.Delay(20);
|
|
|
|
// Only task1 should be running (task2 waiting on the queue slot).
|
|
Assert.Single(order);
|
|
Assert.Equal("1", order[0]);
|
|
|
|
// Release first task.
|
|
gate1.SetResult();
|
|
|
|
// Wait for second task to complete.
|
|
await gate2.Task.WaitAsync(TimeSpan.FromSeconds(5));
|
|
|
|
Assert.Equal(2, order.Count);
|
|
Assert.Equal("2", order[1]);
|
|
|
|
cts.Cancel();
|
|
}
|
|
|
|
[Fact]
|
|
public async Task CancelTask_Triggers_Cancellation()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
|
|
var running = new TaskCompletionSource();
|
|
var cancelled = false;
|
|
|
|
var (service, _) = CreateService(async (_, _, _, _, ct) =>
|
|
{
|
|
running.SetResult();
|
|
try
|
|
{
|
|
await Task.Delay(Timeout.Infinite, ct);
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
cancelled = true;
|
|
throw;
|
|
}
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
});
|
|
|
|
var task = await SeedQueuedTask(listId);
|
|
await service.RunNow(task.Id);
|
|
|
|
await running.Task.WaitAsync(TimeSpan.FromSeconds(5));
|
|
|
|
var result = service.CancelTask(task.Id);
|
|
Assert.True(result);
|
|
|
|
await Task.Delay(200);
|
|
Assert.True(cancelled);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task RunNow_AutoRetries_On_Failure_With_SessionId()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
var task = await SeedQueuedTask(listId);
|
|
|
|
var callCount = 0;
|
|
var (service, fake) = CreateService((prompt, dir, args, onLine, ct) =>
|
|
{
|
|
callCount++;
|
|
if (callCount == 1)
|
|
{
|
|
return Task.FromResult(new RunResult
|
|
{
|
|
ExitCode = 1,
|
|
ErrorMarkdown = "something broke",
|
|
SessionId = "sess-retry-test",
|
|
});
|
|
}
|
|
return Task.FromResult(new RunResult
|
|
{
|
|
ExitCode = 0,
|
|
ResultMarkdown = "fixed it",
|
|
SessionId = "sess-retry-test",
|
|
});
|
|
});
|
|
|
|
await service.StartAsync(CancellationToken.None);
|
|
await service.RunNow(task.Id);
|
|
|
|
// Wait for both runs to complete.
|
|
await Task.Delay(2000);
|
|
|
|
await service.StopAsync(CancellationToken.None);
|
|
|
|
Assert.Equal(2, callCount);
|
|
|
|
var finalTask = await _taskRepo.GetByIdAsync(task.Id);
|
|
Assert.NotNull(finalTask);
|
|
// A standalone task that completes successfully now gates on review.
|
|
Assert.Equal(TaskStatus.WaitingForReview, finalTask.Status);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetActive_Returns_Running_Slots()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
var tcs = new TaskCompletionSource<RunResult>();
|
|
|
|
var (service, _) = CreateService((_, _, _, _, _) => tcs.Task);
|
|
|
|
var task = await SeedQueuedTask(listId);
|
|
await service.RunNow(task.Id);
|
|
|
|
var active = service.GetActive();
|
|
Assert.Single(active);
|
|
Assert.Equal("override", active[0].slot);
|
|
Assert.Equal(task.Id, active[0].taskId);
|
|
|
|
tcs.SetResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" });
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Blocked_UsageGate_Skips_Queue_Refill()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
|
|
var gate = new FakeUsageGate { Decision = new UsageGateDecision(true, "5h-Limit 90% >= 80%") };
|
|
var (service, fake) = CreateService(
|
|
(_, _, _, _, _) => Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" }), gate);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
await Task.Delay(200);
|
|
cts.Cancel();
|
|
|
|
Assert.Equal(0, fake.CallCount);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task UsageGate_Clearing_Resumes_QueueRefill_OnNextTick()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
|
|
var gate = new FakeUsageGate { Decision = new UsageGateDecision(true, "5h-Limit 90% >= 80%") };
|
|
var done = new TaskCompletionSource();
|
|
var (service, fake) = CreateService((_, _, _, _, _) =>
|
|
{
|
|
done.TrySetResult();
|
|
return Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" });
|
|
}, gate);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
await Task.Delay(150);
|
|
Assert.Equal(0, fake.CallCount);
|
|
|
|
// Clear the gate; the 50ms backstop timer in this test's config picks it up.
|
|
gate.Decision = new UsageGateDecision(false, null);
|
|
await done.Task.WaitAsync(TimeSpan.FromSeconds(5));
|
|
cts.Cancel();
|
|
|
|
Assert.Equal(1, fake.CallCount);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Blocked_UsageGate_Does_Not_Cancel_AlreadyRunning_Slot()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
|
|
var running = new TaskCompletionSource();
|
|
var cancelled = false;
|
|
var gate = new FakeUsageGate();
|
|
var (service, _) = CreateService(async (_, _, _, _, ct) =>
|
|
{
|
|
running.SetResult();
|
|
try
|
|
{
|
|
await Task.Delay(Timeout.Infinite, ct);
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
cancelled = true;
|
|
throw;
|
|
}
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
}, gate);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
await running.Task.WaitAsync(TimeSpan.FromSeconds(5));
|
|
|
|
// Block after the slot is already running — several backstop ticks pass.
|
|
gate.Decision = new UsageGateDecision(true, "5h-Limit 90% >= 80%");
|
|
await Task.Delay(200);
|
|
|
|
Assert.False(cancelled);
|
|
cts.Cancel();
|
|
}
|
|
|
|
[Fact]
|
|
public async Task UsageGate_TransitionLogging_FiresOncePerChange()
|
|
{
|
|
var gate = new FakeUsageGate { Decision = new UsageGateDecision(true, "5h-Limit 90% >= 80%") };
|
|
var (service, _) = CreateService(
|
|
(_, _, _, _, _) => Task.FromResult(new RunResult { ExitCode = 0, ResultMarkdown = "ok" }), gate);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
|
|
// Wait for the first backstop tick to log the transition, then keep observing: every
|
|
// further tick (50ms interval) sees the same blocked state and must stay silent. A fixed
|
|
// sleep was flaky here — on a loaded machine no tick fit inside the window at all.
|
|
await AssertStableCountAsync(CountWarnCalls, 1);
|
|
|
|
cts.Cancel();
|
|
}
|
|
|
|
private int CountWarnCalls() => _hubContext.Proxy.Calls
|
|
.Count(c => c.Method == "WorkerLog" && (WorkerLogLevel)c.Args[1]! == WorkerLogLevel.Warn);
|
|
|
|
// Polls until `read()` reaches `expected` (or times out), then waits a further grace period
|
|
// to make sure the count doesn't keep climbing past it — needed because slot fills happen
|
|
// concurrently and a fixed sleep is either flaky (too short) or slow (too long).
|
|
private static async Task AssertStableCountAsync(Func<int> read, int expected)
|
|
{
|
|
var deadline = DateTime.UtcNow.AddSeconds(5);
|
|
while (read() < expected && DateTime.UtcNow < deadline)
|
|
await Task.Delay(20);
|
|
|
|
await Task.Delay(250);
|
|
Assert.Equal(expected, read());
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Throttle_StepsDownEffectiveSlots_BelowConfiguredMax()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
await SeedQueuedTask(listId);
|
|
await SeedQueuedTask(listId);
|
|
|
|
await SetAppSettingsAsync(maxParallel: 3);
|
|
|
|
var usageState = new UsageState();
|
|
usageState.ReportSuccess(new UsageSnapshot(
|
|
new UsageBucket(60, null), new UsageBucket(0, null), Array.Empty<UsageLimitRow>(), DateTime.UtcNow));
|
|
|
|
var startedCount = 0;
|
|
var block = new TaskCompletionSource();
|
|
var (service, _) = CreateService(async (_, _, _, _, _) =>
|
|
{
|
|
Interlocked.Increment(ref startedCount);
|
|
await block.Task;
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
}, usageState: usageState);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
|
|
// 60% is between the soft (50) and hard (65) thresholds — capped at 2 slots even
|
|
// though 3 are configured and 3 tasks are queued.
|
|
await AssertStableCountAsync(() => Volatile.Read(ref startedCount), 2);
|
|
|
|
block.SetResult();
|
|
cts.Cancel();
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Throttle_AtHardThreshold_CapsToOneSlot()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
await SeedQueuedTask(listId);
|
|
|
|
await SetAppSettingsAsync(maxParallel: 3);
|
|
|
|
var usageState = new UsageState();
|
|
usageState.ReportSuccess(new UsageSnapshot(
|
|
new UsageBucket(70, null), new UsageBucket(0, null), Array.Empty<UsageLimitRow>(), DateTime.UtcNow));
|
|
|
|
var startedCount = 0;
|
|
var block = new TaskCompletionSource();
|
|
var (service, _) = CreateService(async (_, _, _, _, _) =>
|
|
{
|
|
Interlocked.Increment(ref startedCount);
|
|
await block.Task;
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
}, usageState: usageState);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
|
|
await AssertStableCountAsync(() => Volatile.Read(ref startedCount), 1);
|
|
|
|
block.SetResult();
|
|
cts.Cancel();
|
|
}
|
|
|
|
[Fact]
|
|
public async Task NoUsageSnapshot_FallsBackToFullConfiguredParallelism()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
await SeedQueuedTask(listId);
|
|
await SeedQueuedTask(listId);
|
|
|
|
await SetAppSettingsAsync(maxParallel: 3);
|
|
|
|
// No snapshot has landed yet (fresh UsageState) — throttle must fail open.
|
|
var startedCount = 0;
|
|
var block = new TaskCompletionSource();
|
|
var (service, _) = CreateService(async (_, _, _, _, _) =>
|
|
{
|
|
Interlocked.Increment(ref startedCount);
|
|
await block.Task;
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
}, usageState: new UsageState());
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
|
|
await AssertStableCountAsync(() => Volatile.Read(ref startedCount), 3);
|
|
|
|
block.SetResult();
|
|
cts.Cancel();
|
|
}
|
|
|
|
[Fact]
|
|
public async Task Throttle_Engaging_Does_Not_Cancel_AlreadyRunning_Slot()
|
|
{
|
|
var listId = await SeedListAsync();
|
|
await SeedQueuedTask(listId);
|
|
|
|
await SetAppSettingsAsync(maxParallel: 3);
|
|
var usageState = new UsageState();
|
|
|
|
var running = new TaskCompletionSource();
|
|
var cancelled = false;
|
|
var (service, _) = CreateService(async (_, _, _, _, ct) =>
|
|
{
|
|
running.SetResult();
|
|
try
|
|
{
|
|
await Task.Delay(Timeout.Infinite, ct);
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
cancelled = true;
|
|
throw;
|
|
}
|
|
return new RunResult { ExitCode = 0, ResultMarkdown = "ok" };
|
|
}, usageState: usageState);
|
|
|
|
using var cts = new CancellationTokenSource();
|
|
await service.StartAsync(cts.Token);
|
|
_waker.Wake();
|
|
await running.Task.WaitAsync(TimeSpan.FromSeconds(5));
|
|
|
|
// Throttle engages hard after the slot is already running — several backstop ticks pass.
|
|
usageState.ReportSuccess(new UsageSnapshot(
|
|
new UsageBucket(70, null), new UsageBucket(0, null), Array.Empty<UsageLimitRow>(), DateTime.UtcNow));
|
|
await Task.Delay(200);
|
|
|
|
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);
|
|
}
|
|
}
|