diff --git a/apps/api/internal/runner/admission.go b/apps/api/internal/runner/admission.go index 662348f..c1cec15 100644 --- a/apps/api/internal/runner/admission.go +++ b/apps/api/internal/runner/admission.go @@ -596,7 +596,7 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) { } continue } - if task.Status != "queued" || task.RiverJobID > 0 { + if task.Status != "queued" { continue } if err := s.dispatchWaitingAsyncTask(ctx, task); err != nil { diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index aeea9a5..f3449a1 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -502,7 +502,11 @@ func (s *Service) enqueueAsyncTaskWithOptions(ctx context.Context, taskID string if err != nil { return err } - defer tx.Rollback(ctx) + defer func() { + rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = tx.Rollback(rollbackCtx) + }() if err := s.enqueueAsyncTaskTx(ctx, tx, taskID, opts); err != nil { return err } @@ -548,18 +552,15 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error { if err != nil { return err } + recovered := 0 for _, item := range items { - if admission, admissionErr := s.store.GetTaskAdmission(ctx, item.ID); admissionErr == nil && admission.Status == "waiting" { - continue - } else if admissionErr != nil && !errors.Is(admissionErr, pgx.ErrNoRows) { - return admissionErr - } if err := s.enqueueAsyncTaskWithOptions(ctx, item.ID, asyncTaskRecoveryInsertOpts(item, time.Now())); err != nil { return err } + recovered++ } - if len(items) > 0 { - s.logger.Info("river async queue recovered persisted tasks", "count", len(items)) + if recovered > 0 { + s.logger.Info("river async queue recovered persisted tasks", "count", recovered) } return nil } diff --git a/apps/api/internal/store/admission_queue.go b/apps/api/internal/store/admission_queue.go index 307c8df..5473471 100644 --- a/apps/api/internal/store/admission_queue.go +++ b/apps/api/internal/store/admission_queue.go @@ -684,11 +684,19 @@ func (s *Store) ListWaitingAsyncAdmissionTaskIDs(ctx context.Context, limit int) limit = 100 } rows, err := s.pool.Query(ctx, ` -SELECT task_id::text -FROM gateway_task_admissions -WHERE status = 'waiting' - AND mode = 'async' -ORDER BY priority ASC, enqueued_at ASC, task_id ASC +SELECT admission.task_id::text +FROM gateway_task_admissions admission +JOIN gateway_tasks task ON task.id = admission.task_id +LEFT JOIN river_job job ON job.id = task.river_job_id +WHERE admission.status = 'waiting' + AND admission.mode = 'async' + AND task.status = 'queued' + AND ( + task.river_job_id IS NULL + OR job.id IS NULL + OR job.state NOT IN ('available', 'pending', 'retryable', 'running', 'scheduled') + ) +ORDER BY admission.priority ASC, admission.enqueued_at ASC, admission.task_id ASC LIMIT $1`, limit) if err != nil { return nil, err diff --git a/apps/api/internal/store/admission_queue_integration_test.go b/apps/api/internal/store/admission_queue_integration_test.go index 1275255..648b8dc 100644 --- a/apps/api/internal/store/admission_queue_integration_test.go +++ b/apps/api/internal/store/admission_queue_integration_test.go @@ -447,6 +447,29 @@ WHERE id = $1::uuid`, queuedAtomicTask.ID).Scan(&riverJobID, &queuedLeases); err if riverJobID != queuedSyntheticRiverJobID || queuedLeases != 0 { t.Fatalf("queued River marker=%d leases=%d, want %d/0", riverJobID, queuedLeases, queuedSyntheticRiverJobID) } + waitingAsyncTaskIDs, err := first.ListWaitingAsyncAdmissionTaskIDs(ctx, 1000) + if err != nil { + t.Fatalf("list waiting async tasks with missing River job: %v", err) + } + waitingAsyncFound := false + for _, taskID := range waitingAsyncTaskIDs { + if taskID == queuedAtomicTask.ID { + waitingAsyncFound = true + break + } + } + if !waitingAsyncFound { + t.Fatal("waiting async task with missing River job was not dispatchable") + } + recoverableTasks, err := first.ListRecoverableAsyncTasks(ctx, 1000) + if err != nil { + t.Fatalf("list recoverable tasks while admission is waiting: %v", err) + } + for _, item := range recoverableTasks { + if item.ID == queuedAtomicTask.ID { + t.Fatal("generic River recovery claimed a task owned by the async admission dispatcher") + } + } result, err = second.TryTaskAdmission(ctx, inputFor(queuedAtomicTask, 100, "")) if err != nil || !result.Admitted || len(result.Leases) != 1 { t.Fatalf("worker-time queued admission result=%+v err=%v", result, err) diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index 5894357..77d52aa 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -947,6 +947,12 @@ WHERE task.async_mode = true OR job.id IS NULL OR job.state NOT IN ('available', 'pending', 'retryable', 'running', 'scheduled') ) + AND NOT EXISTS ( + SELECT 1 + FROM gateway_task_admissions admission + WHERE admission.task_id = task.id + AND admission.status = 'waiting' + ) ORDER BY task.priority ASC, task.created_at ASC LIMIT $1`, limit) if err != nil {