374 lines
15 KiB
C#
374 lines
15 KiB
C#
using ClaudeDo.Data;
|
|
using ClaudeDo.Data.Models;
|
|
using ClaudeDo.Data.Repositories;
|
|
using ClaudeDo.Worker.Config;
|
|
using ClaudeDo.Worker.Hub;
|
|
using ClaudeDo.Worker.Runner;
|
|
using ClaudeDo.Worker.State;
|
|
using ClaudeDo.Worker.Usage;
|
|
using ClaudeDo.Worker.Usage.Interfaces;
|
|
using Microsoft.EntityFrameworkCore;
|
|
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;
|
|
|
|
namespace ClaudeDo.Worker.Queue;
|
|
|
|
public sealed class QueueService : BackgroundService
|
|
{
|
|
private readonly IDbContextFactory<ClaudeDoDbContext> _dbFactory;
|
|
private readonly TaskRunner _runner;
|
|
private readonly WorkerConfig _cfg;
|
|
private readonly ILogger<QueueService> _logger;
|
|
private readonly QueueWaker _waker;
|
|
private readonly IQueuePicker _picker;
|
|
private readonly OverrideSlotService _override;
|
|
private readonly ITaskStateService _state;
|
|
private readonly RunCancellationRegistry _runCancels;
|
|
private readonly IUsageGate _usageGate;
|
|
private readonly UsageState _usageState;
|
|
private readonly HubBroadcaster _broadcaster;
|
|
|
|
private readonly object _lock = new();
|
|
private readonly Dictionary<string, QueueSlotState> _queueSlots = new();
|
|
private bool _usageGateBlocked;
|
|
private int? _lastEffectiveSlots;
|
|
|
|
public QueueService(
|
|
IDbContextFactory<ClaudeDoDbContext> dbFactory,
|
|
TaskRunner runner,
|
|
WorkerConfig cfg,
|
|
ILogger<QueueService> logger,
|
|
QueueWaker waker,
|
|
IQueuePicker picker,
|
|
OverrideSlotService overrideSlot,
|
|
ITaskStateService state,
|
|
RunCancellationRegistry runCancels,
|
|
IUsageGate usageGate,
|
|
UsageState usageState,
|
|
HubBroadcaster broadcaster)
|
|
{
|
|
_dbFactory = dbFactory;
|
|
_runner = runner;
|
|
_cfg = cfg;
|
|
_logger = logger;
|
|
_waker = waker;
|
|
_picker = picker;
|
|
_override = overrideSlot;
|
|
_state = state;
|
|
_runCancels = runCancels;
|
|
_usageGate = usageGate;
|
|
_usageState = usageState;
|
|
_broadcaster = broadcaster;
|
|
}
|
|
|
|
public IReadOnlyList<(string slot, string taskId, DateTime startedAt)> GetActive()
|
|
{
|
|
var list = new List<(string, string, DateTime)>();
|
|
lock (_lock)
|
|
{
|
|
foreach (var slot in _queueSlots.Values)
|
|
list.Add(("queue", slot.TaskId, slot.StartedAt));
|
|
}
|
|
var o = _override.CurrentSlot;
|
|
if (o is not null) list.Add(("override", o.TaskId, o.StartedAt));
|
|
return list;
|
|
}
|
|
|
|
public Task RunNow(string taskId)
|
|
{
|
|
EnsureNotInQueueSlot(taskId);
|
|
return _override.RunNow(taskId);
|
|
}
|
|
|
|
public Task<string> ContinueTask(string taskId, string followUpPrompt)
|
|
{
|
|
EnsureNotInQueueSlot(taskId);
|
|
return _override.ContinueTask(taskId, followUpPrompt);
|
|
}
|
|
|
|
private void EnsureNotInQueueSlot(string taskId)
|
|
{
|
|
lock (_lock)
|
|
{
|
|
if (_queueSlots.ContainsKey(taskId))
|
|
throw new InvalidOperationException("task is already running in queue slot");
|
|
}
|
|
}
|
|
|
|
public bool CancelTask(string taskId)
|
|
{
|
|
if (_override.TryCancel(taskId)) return true;
|
|
|
|
lock (_lock)
|
|
{
|
|
if (_queueSlots.TryGetValue(taskId, out var slot))
|
|
{
|
|
slot.Cts.Cancel();
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
protected override async Task ExecuteAsync(CancellationToken stoppingToken)
|
|
{
|
|
_logger.LogInformation("QueueService started");
|
|
|
|
using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(_cfg.QueueBackstopIntervalMs));
|
|
|
|
while (!stoppingToken.IsCancellationRequested)
|
|
{
|
|
try
|
|
{
|
|
// Wait for wake signal or backstop timer.
|
|
var wakeTask = _waker.WaitAsync(stoppingToken);
|
|
var timerTask = timer.WaitForNextTickAsync(stoppingToken).AsTask();
|
|
|
|
await Task.WhenAny(wakeTask, timerTask);
|
|
|
|
var (_, maxParallel) = await GetSlotCountsAsync(stoppingToken);
|
|
|
|
var gateDecision = await _usageGate.EvaluateAsync(stoppingToken);
|
|
await ReportUsageGateTransitionAsync(gateDecision);
|
|
|
|
// Only queue refill is gated. Runs already in flight (RunNow, ContinueTask,
|
|
// interactive sessions, planning, daily prep) keep going regardless.
|
|
if (!gateDecision.IsBlocked)
|
|
{
|
|
// Fill as many free slots as the limit allows.
|
|
while (!stoppingToken.IsCancellationRequested)
|
|
{
|
|
lock (_lock)
|
|
{
|
|
if (_queueSlots.Count >= maxParallel) break;
|
|
}
|
|
|
|
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)
|
|
{
|
|
_queueSlots[task.Id] = new QueueSlotState { TaskId = task.Id, StartedAt = DateTime.UtcNow, Cts = 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);
|
|
}
|
|
}
|
|
}
|
|
catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
|
|
{
|
|
break;
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
_logger.LogError(ex, "QueueService loop error");
|
|
}
|
|
}
|
|
|
|
_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;
|
|
_usageGateBlocked = decision.IsBlocked;
|
|
|
|
if (decision.IsBlocked)
|
|
{
|
|
_logger.LogInformation("QueueService: usage gate blocking queue refill ({Reason})", decision.Reason);
|
|
await _broadcaster.WorkerLog($"Queue pausiert: {decision.Reason}", WorkerLogLevel.Warn, DateTime.UtcNow);
|
|
}
|
|
else
|
|
{
|
|
var snapshot = _usageState.Snapshot;
|
|
var message = snapshot?.FiveHour is not null && snapshot.SevenDay is not null
|
|
? $"Queue fortgesetzt: 5h {snapshot.FiveHour.Utilization:0}%, 7d {snapshot.SevenDay.Utilization:0}%"
|
|
: "Queue fortgesetzt";
|
|
|
|
_logger.LogInformation("QueueService: usage gate cleared, queue refill resumed");
|
|
await _broadcaster.WorkerLog(message, WorkerLogLevel.Info, DateTime.UtcNow);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Configured parallelism, and that same value stepped down by <see cref="UsageThrottle"/>
|
|
/// ahead of the hard usage gate. A missing snapshot (poll hasn't landed / endpoint
|
|
/// unreachable) fails open to the configured value — a broken usage poll must never stall
|
|
/// the queue. Also called by <c>get_queue_state</c> (External MCP) to surface the throttle
|
|
/// from outside the process.
|
|
/// </summary>
|
|
public async Task<(int Configured, int Effective)> GetSlotCountsAsync(CancellationToken ct)
|
|
{
|
|
int configured;
|
|
UsageThresholds fiveHour, sevenDay;
|
|
try
|
|
{
|
|
using var context = _dbFactory.CreateDbContext();
|
|
var settings = await new AppSettingsRepository(context).GetAsync(ct);
|
|
configured = Math.Max(1, settings.MaxParallelExecutions);
|
|
fiveHour = new UsageThresholds(
|
|
settings.UsageThrottleFiveHourSoftPct, settings.UsageThrottleFiveHourHardPct,
|
|
settings.UsageGateFiveHourPct);
|
|
sevenDay = new UsageThresholds(
|
|
settings.UsageThrottleSevenDaySoftPct, settings.UsageThrottleSevenDayHardPct,
|
|
settings.UsageGateSevenDayPct);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
_logger.LogWarning(ex, "Failed to read max parallel executions; defaulting to 1");
|
|
return (1, 1);
|
|
}
|
|
|
|
var snapshot = _usageState.Snapshot;
|
|
if (snapshot is null || _usageState.LastError is not null)
|
|
{
|
|
_lastEffectiveSlots = configured;
|
|
return (configured, configured);
|
|
}
|
|
|
|
var effective = UsageThrottle.EffectiveSlots(
|
|
configured,
|
|
snapshot.FiveHour?.Utilization, fiveHour,
|
|
snapshot.SevenDay?.Utilization, sevenDay);
|
|
|
|
ReportThrottleTransition(configured, effective, snapshot);
|
|
return (configured, effective);
|
|
}
|
|
|
|
private void ReportThrottleTransition(int configured, int effective, UsageSnapshot snapshot)
|
|
{
|
|
if (_lastEffectiveSlots == effective) return;
|
|
var previous = _lastEffectiveSlots;
|
|
_lastEffectiveSlots = effective;
|
|
|
|
if (previous is null && effective == configured) return; // baseline, nothing to report
|
|
|
|
if (effective < configured)
|
|
{
|
|
_logger.LogInformation(
|
|
"QueueService: usage throttle stepped to {Effective}/{Configured} slots (5h={FiveHour}%, 7d={SevenDay}%)",
|
|
effective, configured, snapshot.FiveHour?.Utilization, snapshot.SevenDay?.Utilization);
|
|
}
|
|
else
|
|
{
|
|
_logger.LogInformation("QueueService: usage throttle cleared, back to {Configured} slots", configured);
|
|
}
|
|
}
|
|
|
|
private async Task RunInSlotAsync(string taskId, CancellationToken ct)
|
|
{
|
|
try
|
|
{
|
|
_logger.LogInformation("Starting task {TaskId} in queue slot", taskId);
|
|
|
|
TaskEntity task;
|
|
using (var context = _dbFactory.CreateDbContext())
|
|
{
|
|
var taskRepo = new TaskRepository(context);
|
|
task = await taskRepo.GetByIdAsync(taskId, ct)
|
|
?? throw new KeyNotFoundException($"Task '{taskId}' not found.");
|
|
}
|
|
|
|
// A task re-queued from review carries reviewer feedback. Resume the prior
|
|
// Claude session with that feedback as the next turn when a session exists;
|
|
// otherwise fall back to a fresh run with the feedback folded into the prompt.
|
|
if (!string.IsNullOrWhiteSpace(task.ReviewFeedback))
|
|
{
|
|
var feedback = task.ReviewFeedback!;
|
|
string? sessionId;
|
|
using (var context = _dbFactory.CreateDbContext())
|
|
sessionId = (await new TaskRunRepository(context).GetLatestByTaskIdAsync(taskId, ct))?.SessionId;
|
|
|
|
if (sessionId is not null)
|
|
{
|
|
await _runner.ContinueAsync(taskId, feedback, "queue", ct, alreadyClaimed: true);
|
|
}
|
|
else
|
|
{
|
|
task.Description = string.IsNullOrWhiteSpace(task.Description)
|
|
? $"Reviewer feedback: {feedback}"
|
|
: $"{task.Description}\n\nReviewer feedback: {feedback}";
|
|
await _runner.RunAsync(task, "queue", ct, alreadyClaimed: true);
|
|
}
|
|
|
|
// Clear the consumed feedback only once the run reached a successful
|
|
// terminal state, so a failed or cancelled run keeps it for a manual retry.
|
|
TaskStatus statusAfter;
|
|
using (var context = _dbFactory.CreateDbContext())
|
|
statusAfter = await context.Tasks.Where(t => t.Id == taskId)
|
|
.Select(t => t.Status).FirstAsync(CancellationToken.None);
|
|
if (statusAfter is TaskStatus.WaitingForReview or TaskStatus.Done)
|
|
await _state.ClearReviewFeedbackAsync(taskId, CancellationToken.None);
|
|
return;
|
|
}
|
|
|
|
await _runner.RunAsync(task, "queue", ct, alreadyClaimed: true);
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
// Cancellation is driven by the cancel path, which already wrote the terminal status.
|
|
// Marking the task Failed here would be a regression (it would stomp Cancelled).
|
|
_logger.LogInformation("Slot runner cancelled for task {TaskId}", taskId);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
_logger.LogError(ex, "Slot runner error for task {TaskId}", taskId);
|
|
|
|
// The picker already committed status='running' before this ran. Without this the
|
|
// task stays Running forever and the UI never hears about it — it keeps showing the
|
|
// pre-claim status because the raw-SQL claim itself never broadcasts.
|
|
try
|
|
{
|
|
await _state.FailAsync(taskId, DateTime.UtcNow,
|
|
$"Slot runner error: {ex.Message}", CancellationToken.None);
|
|
}
|
|
catch (Exception failEx)
|
|
{
|
|
_logger.LogError(failEx, "Could not mark task {TaskId} as failed after a slot error", taskId);
|
|
}
|
|
}
|
|
}
|
|
}
|