diff --git a/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go b/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go index bcb0edf..6b6b9b5 100644 --- a/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go +++ b/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go @@ -96,6 +96,141 @@ SELECT assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs)) } +func TestAsyncAdmissionPoisonTaskDoesNotBlockHealthyFollower(t *testing.T) { + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run async admission resilience tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + applyMigration(t, ctx, databaseURL) + + db, err := store.Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect store: %v", err) + } + defer db.Close() + pool, err := pgxpool.New(ctx, databaseURL) + if err != nil { + t.Fatalf("connect acceptance pool: %v", err) + } + defer pool.Close() + restoreAcceptanceState := isolateHTTPAcceptanceState(t, ctx, pool, false) + defer restoreAcceptanceState() + + serverCtx, cancelServer := context.WithCancel(ctx) + defer cancelServer() + server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{ + AppEnv: "test", + HTTPAddr: ":0", + DatabaseURL: databaseURL, + IdentityMode: "hybrid", + JWTSecret: "test-secret", + BillingEngineMode: "observe", + CORSAllowedOrigin: "*", + AsyncQueueWorkerEnabled: true, + AsyncWorkerHardLimit: 1, + AsyncWorkerInstanceHardLimit: 1, + AsyncWorkerRefreshIntervalSeconds: 1, + AsyncAdmissionMicrobatchSize: 3, + AsyncAdmissionDispatcherEnabled: true, + AsyncAdmissionDispatcherConfigured: true, + }, db, slog.New(slog.NewTextHandler(io.Discard, nil)))) + defer server.Close() + + adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL) + if _, err := pool.Exec(ctx, ` +UPDATE gateway_user_groups +SET rate_limit_policy = '{"rules":[{"metric":"concurrent","limit":256,"leaseTtlSeconds":120}]}'::jsonb +WHERE status = 'active'`); err != nil { + t.Fatalf("raise resilience user-group concurrency: %v", err) + } + suffix := strconv.FormatInt(time.Now().UnixNano(), 10) + model := "admission-poison-" + suffix + platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "admission-poison-"+suffix, "Admission Poison Simulation", 3) + defer updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 3, "disabled") + createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "inherit", 0, 120) + + waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 1, 15*time.Second) + taskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 3, 10*time.Second) + waitForTaskCount(t, ctx, pool, taskIDs, "running", 1, 10*time.Second) + + var poisonTaskID, followerTaskID string + if err := pool.QueryRow(ctx, ` +SELECT (array_agg(admission.task_id::text ORDER BY admission.priority, admission.enqueued_at, admission.task_id))[1], + (array_agg(admission.task_id::text ORDER BY admission.priority, admission.enqueued_at, admission.task_id))[2] +FROM gateway_task_admissions admission +JOIN gateway_tasks task ON task.id = admission.task_id +WHERE task.id = ANY($1::uuid[]) + AND task.status = 'queued' + AND admission.status = 'waiting'`, taskIDs).Scan(&poisonTaskID, &followerTaskID); err != nil { + t.Fatalf("select waiting poison and follower tasks: %v", err) + } + var prematureReselect int + if err := pool.QueryRow(ctx, ` +SELECT count(*) +FROM gateway_task_admissions +WHERE task_id = ANY($1::uuid[]) + AND status = 'waiting' + AND reselect_requested_at IS NOT NULL`, taskIDs).Scan(&prematureReselect); err != nil { + t.Fatalf("count global-capacity reselections: %v", err) + } + if prematureReselect != 0 { + t.Fatalf("global worker saturation marked %d tasks for reselection", prematureReselect) + } + if _, err := pool.Exec(ctx, ` +UPDATE gateway_tasks +SET model = $2, + requested_model = $2, + updated_at = now() +WHERE id = $1::uuid`, poisonTaskID, "missing-model-"+suffix); err != nil { + t.Fatalf("poison queued task model: %v", err) + } + if _, err := pool.Exec(ctx, ` +UPDATE gateway_task_admissions +SET reselect_requested_at = now(), + dispatch_next_at = now(), + updated_at = now() +WHERE task_id = $1::uuid`, poisonTaskID); err != nil { + t.Fatalf("poison queued admission binding: %v", err) + } + + waitForTaskCount(t, ctx, pool, []string{poisonTaskID}, "failed", 1, 20*time.Second) + waitForTaskCount(t, ctx, pool, []string{followerTaskID}, "succeeded", 1, 30*time.Second) + var poisonCode string + var poisonAttempts, poisonAdmissions, poisonLeases, poisonFailedEvents int + var poisonRiverJobID int64 + if err := pool.QueryRow(ctx, ` +SELECT COALESCE(task.error_code, ''), COALESCE(task.river_job_id, 0), + (SELECT count(*) FROM gateway_task_attempts attempt WHERE attempt.task_id = task.id), + (SELECT count(*) FROM gateway_task_admissions admission WHERE admission.task_id = task.id), + (SELECT count(*) FROM gateway_concurrency_leases lease WHERE lease.task_id = task.id AND lease.released_at IS NULL), + (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.failed') +FROM gateway_tasks task +WHERE task.id = $1::uuid`, poisonTaskID).Scan( + &poisonCode, + &poisonRiverJobID, + &poisonAttempts, + &poisonAdmissions, + &poisonLeases, + &poisonFailedEvents, + ); err != nil { + t.Fatalf("read poison task terminal state: %v", err) + } + if poisonCode != "no_model_candidate" || poisonRiverJobID != 0 || poisonAttempts != 0 || + poisonAdmissions != 0 || poisonLeases != 0 || poisonFailedEvents != 1 { + t.Fatalf( + "poison state code=%s river=%v attempts=%d admissions=%d leases=%d events=%d", + poisonCode, + poisonRiverJobID, + poisonAttempts, + poisonAdmissions, + poisonLeases, + poisonFailedEvents, + ) + } +} + func TestAsyncWorkerThousandConcurrentSchedulingAcceptance(t *testing.T) { databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) if databaseURL == "" { @@ -537,7 +672,24 @@ func waitForTaskCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tas } time.Sleep(100 * time.Millisecond) } - t.Fatalf("task status %s count did not reach %d within %s", status, count, timeout) + rows, err := pool.Query(ctx, ` +SELECT task.status, COALESCE(task.error_code, ''), count(*) +FROM gateway_tasks task +WHERE task.id = ANY($1::uuid[]) +GROUP BY task.status, COALESCE(task.error_code, '') +ORDER BY task.status, COALESCE(task.error_code, '')`, taskIDs) + diagnostics := make([]string, 0) + if err == nil { + defer rows.Close() + for rows.Next() { + var taskStatus, code string + var taskCount int + if scanErr := rows.Scan(&taskStatus, &code, &taskCount); scanErr == nil { + diagnostics = append(diagnostics, fmt.Sprintf("%s/%s=%d", taskStatus, code, taskCount)) + } + } + } + t.Fatalf("task status %s count did not reach %d within %s; observed=%v", status, count, timeout, diagnostics) } func activeLeaseExpiry(t *testing.T, ctx context.Context, pool *pgxpool.Pool, modelID string, wantCount int) time.Time { diff --git a/apps/api/internal/runner/admission.go b/apps/api/internal/runner/admission.go index 5299b49..4cf0e93 100644 --- a/apps/api/internal/runner/admission.go +++ b/apps/api/internal/runner/admission.go @@ -745,17 +745,19 @@ func admissionInputFromSnapshot(admission store.TaskAdmission, workerCapacity in } } -func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []store.TaskAdmission) (bool, error) { +func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []store.TaskAdmission) (bool, bool, error) { if len(admissions) == 0 { - return false, nil + return false, false, nil } workerCapacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx) if err != nil { - return false, err + return false, false, err } if workerCapacity <= 0 { - return true, nil + s.observeAsyncAdmissionDispatch("global_wait") + return true, false, nil } + progressed := false microbatchSize := s.cfg.AsyncAdmissionMicrobatchSize if microbatchSize <= 0 { microbatchSize = 1 @@ -765,7 +767,7 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []st inputs := make([]store.TaskAdmissionInput, 0, end-start) tasksByID := make(map[string]store.GatewayTask, end-start) for _, admission := range admissions[start:end] { - if admission.ReselectRequestedAt.IsZero() && len(admission.Scopes) > 0 { + if admission.ReselectRequestedAt.IsZero() && len(admission.Scopes) > 0 && !admission.SnapshotInvalid { inputs = append(inputs, admissionInputFromSnapshot(admission, workerCapacity)) tasksByID[admission.TaskID] = store.GatewayTask{ID: admission.TaskID} continue @@ -773,78 +775,411 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []st task, loadErr := s.store.GetTask(ctx, admission.TaskID) if loadErr != nil { if errors.Is(loadErr, pgx.ErrNoRows) { + progressed = true continue } - return false, loadErr + return false, progressed, loadErr } plan, planErr := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, authUserFromTask(task), &admission) if planErr != nil { - return false, planErr + handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, task, admission, planErr) + if handleErr != nil { + return false, progressed, handleErr + } + progressed = progressed || handled + continue } if !plan.Eligible { - if enqueueErr := s.EnqueueAsyncTask(ctx, task); enqueueErr != nil { - return false, enqueueErr + // This task already owns a durable admission row from its previous + // binding. Promote that row and insert the River job in the same + // transaction even when the replacement candidate has no business + // concurrency scope; direct enqueue would strand a waiting row and + // make every dispatcher enqueue it forever. + if plan.Candidate.PlatformID == "" || plan.Candidate.PlatformModelID == "" { + handled, handleErr := s.handleAsyncAdmissionTaskError( + ctx, + task, + admission, + errors.New("waiting admission rebuilt without a platform model candidate"), + ) + if handleErr != nil { + return false, progressed, handleErr + } + progressed = progressed || handled + continue } - continue } input := taskAdmissionInput(task, plan, "") if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil && !errors.Is(rebindErr, pgx.ErrNoRows) { - return false, rebindErr + handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, task, admission, rebindErr) + if handleErr != nil { + return false, progressed, handleErr + } + progressed = progressed || handled + continue } + progressed = true inputs = append(inputs, input) tasksByID[task.ID] = task } if len(inputs) == 0 { continue } - outcomes, err := s.store.TryTaskAdmissionAtomicBatchWithAdmittedHook( + enqueueHook := func(tx pgx.Tx, input store.TaskAdmissionInput) error { + task, ok := tasksByID[input.TaskID] + if !ok { + return fmt.Errorf("async admission batch task %s was not prepared", input.TaskID) + } + return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task)) + } + outcomes, batchErr := s.store.TryTaskAdmissionAtomicBatchWithAdmittedHook( ctx, inputs, - func(tx pgx.Tx, input store.TaskAdmissionInput) error { - task, ok := tasksByID[input.TaskID] - if !ok { - return fmt.Errorf("async admission batch task %s was not prepared", input.TaskID) - } - return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task)) - }, + enqueueHook, ) - if err != nil { - return false, err - } - for _, outcome := range outcomes { - s.observeTaskAdmissionAttempt(outcome.Result, outcome.Err) - if outcome.Err != nil { - if errors.Is(outcome.Err, store.ErrTaskExecutionFinished) || - errors.Is(outcome.Err, store.ErrQueueTimeout) { + if batchErr != nil { + if isAsyncAdmissionSystemError(batchErr) { + return false, progressed, batchErr + } + failedTaskID := asyncAdmissionFailedTaskID(outcomes) + if failedTaskID == "" { + return false, progressed, batchErr + } + failedTask, loadErr := s.asyncAdmissionDispatchTask(ctx, tasksByID[failedTaskID]) + if loadErr != nil { + return false, progressed, loadErr + } + failedAdmission := admissionByTaskID(admissions[start:end], failedTaskID) + handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, failedTask, failedAdmission, batchErr) + if handleErr != nil { + return false, progressed, handleErr + } + progressed = progressed || handled + + for _, input := range inputs { + if input.TaskID == failedTaskID { continue } - return false, outcome.Err + result, oneErr := s.store.TryTaskAdmissionWithAdmittedHook(ctx, input, func(tx pgx.Tx) error { + return enqueueHook(tx, input) + }) + globalSaturated, outcomeProgress, outcomeErr := s.handleAsyncAdmissionOutcome( + ctx, + store.TaskAdmissionBatchOutcome{TaskID: input.TaskID, Result: result, Err: oneErr}, + input, + tasksByID, + admissionByTaskID(admissions[start:end], input.TaskID), + ) + if outcomeErr != nil { + return false, progressed, outcomeErr + } + progressed = progressed || outcomeProgress + if globalSaturated { + return true, progressed, nil + } } - if !outcome.Result.Admitted { - s.observeCandidateRouting("quota_race_rotated") - platformModelID := outcome.Result.Admission.PlatformModelID - if platformModelID == "" { - for _, input := range inputs { - if input.TaskID == outcome.TaskID { - platformModelID = input.PlatformModelID - break - } - } - } - if platformModelID != "" { - marked, markErr := s.store.RequestWaitingTaskAdmissionReselect(ctx, platformModelID) - if markErr != nil { - return false, markErr - } - if marked > 0 { - s.observeTaskAdmission("candidate_reselect_requested") - } - } - return true, nil + continue + } + for _, outcome := range outcomes { + input := inputByTaskID(inputs, outcome.TaskID) + globalSaturated, outcomeProgress, outcomeErr := s.handleAsyncAdmissionOutcome( + ctx, + outcome, + input, + tasksByID, + admissionByTaskID(admissions[start:end], outcome.TaskID), + ) + if outcomeErr != nil { + return false, progressed, outcomeErr + } + progressed = progressed || outcomeProgress + if globalSaturated { + return true, progressed, nil } } } - return false, nil + return false, progressed, nil +} + +func (s *Service) handleAsyncAdmissionOutcome( + ctx context.Context, + outcome store.TaskAdmissionBatchOutcome, + input store.TaskAdmissionInput, + tasksByID map[string]store.GatewayTask, + admission store.TaskAdmission, +) (bool, bool, error) { + s.observeTaskAdmissionAttempt(outcome.Result, outcome.Err) + if outcome.Err != nil { + if errors.Is(outcome.Err, store.ErrTaskExecutionFinished) || errors.Is(outcome.Err, store.ErrQueueTimeout) { + return false, true, nil + } + if isAsyncAdmissionSystemError(outcome.Err) { + return false, false, outcome.Err + } + task, err := s.asyncAdmissionDispatchTask(ctx, tasksByID[outcome.TaskID]) + if err != nil { + return false, false, err + } + handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, task, admission, outcome.Err) + return false, handled, handleErr + } + if outcome.Result.Admitted { + s.observeAsyncAdmissionDispatch("admitted") + return false, true, nil + } + + globalBlocked := false + groupBlocked := false + platformSaturated := false + for _, blocker := range outcome.Result.Blockers { + switch blocker.ScopeType { + case "worker_capacity": + if blocker.Reason == "saturated" { + globalBlocked = true + } + case "user_group": + groupBlocked = true + case "platform_model": + if blocker.Reason == "saturated" { + platformSaturated = true + } + } + } + if globalBlocked { + s.observeAsyncAdmissionDispatch("global_wait") + return true, false, nil + } + if groupBlocked { + s.observeAsyncAdmissionDispatch("group_wait") + return false, false, nil + } + if platformSaturated { + platformModelID := outcome.Result.Admission.PlatformModelID + if platformModelID == "" { + platformModelID = input.PlatformModelID + } + _, marked, err := s.store.RequestTaskAdmissionReselect(ctx, outcome.TaskID, platformModelID) + if err != nil { + return false, false, err + } + s.observeAsyncAdmissionDispatch("platform_wait") + if marked { + s.observeTaskAdmission("candidate_reselect_requested") + s.observeAsyncAdmissionDispatch("reselect") + return false, true, nil + } + return false, false, nil + } + if len(outcome.Result.Blockers) > 0 { + s.observeAsyncAdmissionDispatch("platform_wait") + return false, false, nil + } + + task, err := s.asyncAdmissionDispatchTask(ctx, tasksByID[outcome.TaskID]) + if err != nil { + return false, false, err + } + handled, handleErr := s.handleAsyncAdmissionTaskError( + ctx, + task, + admission, + errors.New("task admission returned neither admitted nor a blocker"), + ) + return false, handled, handleErr +} + +func (s *Service) handleAsyncAdmissionTaskError( + ctx context.Context, + task store.GatewayTask, + admission store.TaskAdmission, + cause error, +) (bool, error) { + if isAsyncAdmissionSystemError(cause) { + return false, cause + } + now := time.Now() + terminal := asyncAdmissionTaskErrorTerminal(cause) + code := asyncAdmissionTaskErrorCode(cause) + if !admission.WaitDeadlineAt.IsZero() && !admission.WaitDeadlineAt.After(now) { + terminal = true + if !errors.Is(cause, store.ErrNoModelCandidate) { + code = "admission_dispatch_failed" + } + } + if terminal { + callbackURL, callbackErr := s.taskCallbackURL(ctx, task.ID) + if callbackErr != nil { + return false, callbackErr + } + _, failErr := s.store.FailQueuedTaskWithCallback( + context.WithoutCancel(ctx), + task.ID, + code, + cause.Error(), + callbackURL, + task.RunMode == "simulation", + ) + if errors.Is(failErr, store.ErrTaskExecutionFinished) { + return true, nil + } + if failErr != nil { + return false, failErr + } + s.observeAsyncAdmissionDispatch("terminal") + if s.logger != nil { + s.logger.Info("terminally failed waiting async admission", + "taskID", task.ID, + "error_category", code, + ) + } + return true, nil + } + + retryAt := asyncAdmissionRetryAt(now, task.ID, admission.DispatchFailureCount, cause) + updated, changed, deferErr := s.store.DeferWaitingTaskAdmission( + ctx, + task.ID, + admission.PlatformModelID, + retryAt, + "deferred", + code, + ) + if deferErr != nil { + return false, deferErr + } + if !changed { + return true, nil + } + s.observeAsyncAdmissionDispatch("deferred") + s.observeAsyncAdmissionDispatch("task_error") + if s.logger != nil && (updated.DispatchFailureCount == 1 || updated.DispatchFailureCount == 5 || updated.DispatchFailureCount == 7 || updated.DispatchFailureCount%10 == 0) { + s.logger.Warn("deferred task-specific async admission failure", + "taskID", task.ID, + "failureCount", updated.DispatchFailureCount, + "nextDispatchAt", updated.DispatchNextAt.UTC().Format(time.RFC3339), + "error_category", code, + ) + } + return true, nil +} + +func asyncAdmissionTaskErrorTerminal(err error) bool { + if errors.Is(err, store.ErrNoModelCandidate) { + return store.ModelCandidateRetryAfter(err) <= 0 && store.ModelCandidateRecoveryAt(err).IsZero() + } + if errors.Is(err, store.ErrInvalidPlatformModelConfiguration) { + return true + } + if errors.Is(err, store.ErrRateLimited) { + return !store.RateLimitRetryable(err) + } + var clientErr *clients.ClientError + return errors.As(err, &clientErr) && !clientErr.Retryable +} + +func asyncAdmissionTaskErrorCode(err error) string { + if errors.Is(err, store.ErrNoModelCandidate) { + return store.ModelCandidateErrorCode(err) + } + return clients.ErrorCode(err) +} + +func asyncAdmissionRetryAt(now time.Time, taskID string, failureCount int, err error) time.Time { + if recoveryAt := store.ModelCandidateRecoveryAt(err); recoveryAt.After(now) { + return recoveryAt + } + if delay := store.ModelCandidateRetryAfter(err); delay > 0 { + return now.Add(delay) + } + if delay := store.RateLimitRetryAfter(err); delay > 0 { + return now.Add(delay) + } + delay := time.Second + for index := 0; index < failureCount && delay < time.Minute; index++ { + delay *= 2 + } + if delay > time.Minute { + delay = time.Minute + } + hasher := fnv.New32a() + _, _ = hasher.Write([]byte(taskID)) + jitterPermille := int64(hasher.Sum32()%401) - 200 + delay += time.Duration(int64(delay) * jitterPermille / 1000) + if delay < time.Second { + delay = time.Second + } + if delay > time.Minute { + delay = time.Minute + } + return now.Add(delay) +} + +func isAsyncAdmissionSystemError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || store.IsPostgresUnavailable(err) { + return true + } + var sqlState interface{ SQLState() string } + if !errors.As(err, &sqlState) { + return false + } + code := strings.TrimSpace(sqlState.SQLState()) + return strings.HasPrefix(code, "08") || // connection exception + strings.HasPrefix(code, "25") || // invalid transaction state + strings.HasPrefix(code, "40") || // transaction rollback/deadlock + strings.HasPrefix(code, "53") || // insufficient resources + strings.HasPrefix(code, "58") || // system error + code == "55P03" || // lock not available + code == "57014" || // query cancelled/statement timeout + code == "57P01" || code == "57P02" || code == "57P03" +} + +func asyncAdmissionFailedTaskID(outcomes []store.TaskAdmissionBatchOutcome) string { + for index := len(outcomes) - 1; index >= 0; index-- { + if outcomes[index].Err != nil { + return outcomes[index].TaskID + } + } + return "" +} + +func inputByTaskID(inputs []store.TaskAdmissionInput, taskID string) store.TaskAdmissionInput { + for _, input := range inputs { + if input.TaskID == taskID { + return input + } + } + return store.TaskAdmissionInput{TaskID: taskID} +} + +func admissionByTaskID(admissions []store.TaskAdmission, taskID string) store.TaskAdmission { + for _, admission := range admissions { + if admission.TaskID == taskID { + return admission + } + } + return store.TaskAdmission{TaskID: taskID} +} + +func (s *Service) asyncAdmissionDispatchTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) { + if task.ID == "" { + return store.GatewayTask{}, pgx.ErrNoRows + } + if task.Kind != "" { + return task, nil + } + return s.store.GetTask(ctx, task.ID) +} + +func (s *Service) observeAsyncAdmissionDispatch(outcome string) { + observer, ok := s.billingMetrics.(interface { + ObserveAsyncAdmissionDispatch(string) + }) + if ok { + observer.ObserveAsyncAdmissionDispatch(outcome) + } } func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskID string) bool { @@ -853,7 +1188,16 @@ func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskI } cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) defer cancel() - if _, changed, err := s.store.CancelQueuedTask(cleanupCtx, taskID, "client disconnected before upstream submission"); err != nil { + callbackURL, callbackErr := s.taskCallbackURL(cleanupCtx, taskID) + if callbackErr != nil && s.logger != nil { + s.logger.Warn("resolve disconnected task callback failed", "taskID", taskID, "error", callbackErr) + } + if _, changed, err := s.store.CancelQueuedTaskWithCallback( + cleanupCtx, + taskID, + "client disconnected before upstream submission", + callbackURL, + ); err != nil { if s.logger != nil { s.logger.Warn("cancel disconnected asynchronous task failed", "taskID", taskID, "error", err) } @@ -866,6 +1210,8 @@ func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskI func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) { ticker := time.NewTicker(time.Second) defer ticker.Stop() + systemBackoff := time.Second + var systemRetryAt time.Time for { select { case <-ctx.Done(): @@ -873,21 +1219,34 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) { case <-s.asyncAdmissionWake: case <-ticker.C: } + if time.Now().Before(systemRetryAt) { + continue + } for { batchLimit := asyncAdmissionBatchLimit(s.cfg.AsyncWorkerHardLimit) admissions, err := s.store.ListWaitingAsyncAdmissions(ctx, batchLimit) if err != nil { s.logger.Warn("list waiting async admissions failed", "error", err) + s.observeAsyncAdmissionDispatch("system_error") + systemRetryAt = time.Now().Add(systemBackoff) + systemBackoff = min(systemBackoff*2, 30*time.Second) break } if len(admissions) == 0 { + systemBackoff = time.Second + systemRetryAt = time.Time{} break } - saturated, err := s.dispatchWaitingAsyncTasks(ctx, admissions) + saturated, progressed, err := s.dispatchWaitingAsyncTasks(ctx, admissions) if err != nil { s.logger.Warn("dispatch waiting async admission batch failed", "tasks", len(admissions), "error", err) + s.observeAsyncAdmissionDispatch("system_error") + systemRetryAt = time.Now().Add(systemBackoff) + systemBackoff = min(systemBackoff*2, 30*time.Second) break } + systemBackoff = time.Second + systemRetryAt = time.Time{} if saturated { // Every asynchronous task shares the global worker-capacity FIFO // scope. Once its current head cannot be admitted, later tasks @@ -895,6 +1254,9 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) { s.observeTaskAdmission("dispatch_saturated") break } + if !progressed { + break + } // PostgreSQL notifications are wake-up hints and collapse while a // batch is being admitted. Keep filling consecutive batches until // the global execution scope is actually saturated so a large burst @@ -919,7 +1281,11 @@ func (s *Service) reapExpiredTaskAdmissions(ctx context.Context) { return case <-ticker.C: } - result, err := s.store.ReapExpiredTaskAdmissions(ctx, 500) + defaultCallbackURL := "" + if s.cfg.TaskProgressCallbackEnabled { + defaultCallbackURL = s.cfg.TaskProgressCallbackURL + } + result, err := s.store.ReapExpiredTaskAdmissions(ctx, 500, defaultCallbackURL) if err != nil { if s.logger != nil { s.logger.Warn("reap expired task admissions failed", "error", err) diff --git a/apps/api/internal/runner/admission_test.go b/apps/api/internal/runner/admission_test.go index edef8c1..7ad4ca8 100644 --- a/apps/api/internal/runner/admission_test.go +++ b/apps/api/internal/runner/admission_test.go @@ -1,12 +1,84 @@ package runner import ( + "errors" "testing" "time" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" + "github.com/jackc/pgx/v5/pgconn" ) +func TestAsyncAdmissionTaskErrorDisposition(t *testing.T) { + for _, testCase := range []struct { + name string + err error + terminal bool + code string + }{ + {name: "bare no candidate", err: store.ErrNoModelCandidate, terminal: true, code: "no_model_candidate"}, + {name: "cooldown", err: &store.ModelCandidateUnavailableError{Code: "model_cooling_down", Message: "cooling", RetryAfter: time.Minute}, code: "model_cooling_down"}, + {name: "retryable rate limit", err: &store.RateLimitExceededError{Retryable: true}, code: "gateway_rate_limited"}, + {name: "invalid parameter", err: &clients.ClientError{Code: "invalid_parameter", Retryable: false}, terminal: true, code: "invalid_parameter"}, + {name: "retryable client", err: &clients.ClientError{Code: "network", Retryable: true}, code: "network"}, + {name: "unknown", err: errors.New("unknown dispatcher failure"), code: "client_error"}, + } { + t.Run(testCase.name, func(t *testing.T) { + if got := asyncAdmissionTaskErrorTerminal(testCase.err); got != testCase.terminal { + t.Fatalf("terminal=%v, want %v", got, testCase.terminal) + } + if got := asyncAdmissionTaskErrorCode(testCase.err); got != testCase.code { + t.Fatalf("code=%q, want %q", got, testCase.code) + } + }) + } +} + +func TestAsyncAdmissionSystemErrorClassification(t *testing.T) { + for _, testCase := range []struct { + name string + code string + system bool + }{ + {name: "connection", code: "08006", system: true}, + {name: "deadlock", code: "40P01", system: true}, + {name: "resource exhaustion", code: "53300", system: true}, + {name: "task data violation", code: "22023", system: false}, + {name: "task unique violation", code: "23505", system: false}, + } { + t.Run(testCase.name, func(t *testing.T) { + err := &pgconn.PgError{Code: testCase.code} + if got := isAsyncAdmissionSystemError(err); got != testCase.system { + t.Fatalf("SQLSTATE %s system=%v, want %v", testCase.code, got, testCase.system) + } + }) + } +} + +func TestAsyncAdmissionRetryAtUsesRecoveryAndBoundedBackoff(t *testing.T) { + now := time.Unix(1000, 0) + recoveryAt := now.Add(7 * time.Minute) + if got := asyncAdmissionRetryAt(now, "task-recovery", 9, &store.ModelCandidateUnavailableError{ + Message: "cooling", + RetryAfter: time.Minute, + RecoveryAt: recoveryAt, + }); !got.Equal(recoveryAt) { + t.Fatalf("recovery retry=%s, want %s", got, recoveryAt) + } + previous := time.Duration(0) + for failureCount := 0; failureCount < 10; failureCount++ { + delay := asyncAdmissionRetryAt(now, "stable-task", failureCount, errors.New("boom")).Sub(now) + if delay < time.Second || delay > time.Minute { + t.Fatalf("failure %d delay=%s outside jittered 1s..60s bounds", failureCount, delay) + } + if delay < previous { + t.Fatalf("failure %d delay=%s decreased from %s", failureCount, delay, previous) + } + previous = delay + } +} + func TestAdmissionInputFromSnapshotRefreshesWorkerCapacity(t *testing.T) { admission := store.TaskAdmission{ TaskID: "task-1", diff --git a/apps/api/internal/securityevents/metrics.go b/apps/api/internal/securityevents/metrics.go index c2cd446..d474e70 100644 --- a/apps/api/internal/securityevents/metrics.go +++ b/apps/api/internal/securityevents/metrics.go @@ -80,6 +80,15 @@ type Metrics struct { taskAdmissionCancelled atomic.Uint64 taskAdmissionExpired atomic.Uint64 taskAdmissionMigrated atomic.Uint64 + asyncAdmissionAdmitted atomic.Uint64 + asyncAdmissionGlobalWait atomic.Uint64 + asyncAdmissionPlatformWait atomic.Uint64 + asyncAdmissionGroupWait atomic.Uint64 + asyncAdmissionReselect atomic.Uint64 + asyncAdmissionDeferred atomic.Uint64 + asyncAdmissionTerminal atomic.Uint64 + asyncAdmissionTaskError atomic.Uint64 + asyncAdmissionSystemError atomic.Uint64 candidateNormalRotation atomic.Uint64 candidateFullAvoided atomic.Uint64 candidateQuotaRaceRotated atomic.Uint64 @@ -260,6 +269,29 @@ func (m *Metrics) ObserveTaskAdmission(event string) { } } +func (m *Metrics) ObserveAsyncAdmissionDispatch(outcome string) { + switch outcome { + case "admitted": + m.asyncAdmissionAdmitted.Add(1) + case "global_wait": + m.asyncAdmissionGlobalWait.Add(1) + case "platform_wait": + m.asyncAdmissionPlatformWait.Add(1) + case "group_wait": + m.asyncAdmissionGroupWait.Add(1) + case "reselect": + m.asyncAdmissionReselect.Add(1) + case "deferred": + m.asyncAdmissionDeferred.Add(1) + case "terminal": + m.asyncAdmissionTerminal.Add(1) + case "task_error": + m.asyncAdmissionTaskError.Add(1) + case "system_error": + m.asyncAdmissionSystemError.Add(1) + } +} + func (m *Metrics) ObserveCandidateRouting(reason string) { switch reason { case "normal_rotation": @@ -568,7 +600,10 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str plainGauge(w, "easyai_gateway_task_admission_queue_depth", "Current persistent non-text task admission queue depth.", int64(admission.QueueDepth)) plainGauge(w, "easyai_gateway_task_admission_waiting_sync", "Current synchronous admission waiters.", int64(admission.WaitingSync)) plainGauge(w, "easyai_gateway_task_admission_waiting_async", "Current asynchronous admission waiters.", int64(admission.WaitingAsync)) + plainGauge(w, "easyai_gateway_task_admission_deferred_async", "Asynchronous admissions quarantined until a future dispatch attempt.", int64(admission.DeferredAsync)) + plainGauge(w, "easyai_gateway_task_admission_pending_reselect", "Waiting asynchronous admissions marked for candidate reselection.", int64(admission.PendingReselect)) plainFloatGauge(w, "easyai_gateway_task_admission_oldest_wait_seconds", "Age of the oldest persistent admission waiter.", admission.OldestWaitSeconds) + plainFloatGauge(w, "easyai_gateway_task_admission_oldest_dispatchable_wait_seconds", "Age of the oldest asynchronous admission currently eligible for dispatch.", admission.OldestDispatchableWaitSeconds) plainGauge(w, "easyai_gateway_task_admission_expired_waiter_backlog", "Expired synchronous waiter leases awaiting reclamation.", int64(admission.ExpiredWaiterBacklog)) plainGauge(w, "easyai_gateway_task_admission_expired_deadline_backlog", "Expired queue deadlines awaiting reclamation.", int64(admission.ExpiredDeadlineBacklog)) outcomeCounters(w, "easyai_gateway_task_admissions_total", "Task admission transitions by bounded outcome.", []outcomeValue{ @@ -579,6 +614,17 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str {"expired", m.taskAdmissionExpired.Load()}, {"candidate_migrated", m.taskAdmissionMigrated.Load()}, }) + outcomeCounters(w, "easyai_gateway_async_admission_dispatch_total", "Asynchronous admission dispatcher decisions by bounded outcome.", []outcomeValue{ + {"admitted", m.asyncAdmissionAdmitted.Load()}, + {"global_wait", m.asyncAdmissionGlobalWait.Load()}, + {"platform_wait", m.asyncAdmissionPlatformWait.Load()}, + {"group_wait", m.asyncAdmissionGroupWait.Load()}, + {"reselect", m.asyncAdmissionReselect.Load()}, + {"deferred", m.asyncAdmissionDeferred.Load()}, + {"terminal", m.asyncAdmissionTerminal.Load()}, + {"task_error", m.asyncAdmissionTaskError.Load()}, + {"system_error", m.asyncAdmissionSystemError.Load()}, + }) taskAdmissionWaitHistogram(w, m) }) } diff --git a/apps/api/internal/securityevents/metrics_test.go b/apps/api/internal/securityevents/metrics_test.go index cfbc6c4..b551de8 100644 --- a/apps/api/internal/securityevents/metrics_test.go +++ b/apps/api/internal/securityevents/metrics_test.go @@ -62,6 +62,13 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { metrics.ObserveConcurrencyLeaseRenewal("lost") metrics.ObserveTaskEventSkip("duplicate") metrics.ObserveTaskEventSkip("budget_exceeded") + metrics.ObserveAsyncAdmissionDispatch("admitted") + metrics.ObserveAsyncAdmissionDispatch("global_wait") + metrics.ObserveAsyncAdmissionDispatch("reselect") + metrics.ObserveAsyncAdmissionDispatch("deferred") + metrics.ObserveAsyncAdmissionDispatch("terminal") + metrics.ObserveAsyncAdmissionDispatch("task_error") + metrics.ObserveAsyncAdmissionDispatch("system_error") recorder := httptest.NewRecorder() metrics.Handler(metricsSnapshot{ @@ -116,6 +123,13 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { `easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`, `easyai_gateway_task_events_skipped_total{outcome="duplicate"} 1`, `easyai_gateway_task_events_skipped_total{outcome="budget_exceeded"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="admitted"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="global_wait"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="reselect"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="deferred"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="terminal"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="task_error"} 1`, + `easyai_gateway_async_admission_dispatch_total{outcome="system_error"} 1`, `easyai_gateway_postgres_pool_max_connections 32`, `easyai_gateway_postgres_pool_total_connections 18`, `easyai_gateway_postgres_pool_acquired_connections 12`, diff --git a/apps/api/internal/store/admission_queue.go b/apps/api/internal/store/admission_queue.go index aed5945..1797449 100644 --- a/apps/api/internal/store/admission_queue.go +++ b/apps/api/internal/store/admission_queue.go @@ -65,21 +65,34 @@ type TaskAdmissionInput struct { } type TaskAdmission struct { - TaskID string - PlatformID string - PlatformModelID string - UserGroupID string - QueueKey string - Mode string - Status string - Priority int - EnqueuedAt time.Time - WaitDeadlineAt time.Time - AdmittedAt time.Time - WaiterID string - WaiterLeaseExpires time.Time - Scopes []AdmissionScope - ReselectRequestedAt time.Time + TaskID string + PlatformID string + PlatformModelID string + UserGroupID string + QueueKey string + Mode string + Status string + Priority int + EnqueuedAt time.Time + WaitDeadlineAt time.Time + AdmittedAt time.Time + WaiterID string + WaiterLeaseExpires time.Time + Scopes []AdmissionScope + ReselectRequestedAt time.Time + DispatchNextAt time.Time + DispatchFailureCount int + DispatchLastOutcome string + DispatchLastErrorCode string + DispatchLastErrorAt time.Time + SnapshotInvalid bool +} + +type TaskAdmissionBlocker struct { + Reason string + ScopeType string + ScopeKey string + RetryAt time.Time } type TaskAdmissionResult struct { @@ -87,6 +100,7 @@ type TaskAdmissionResult struct { Admitted bool NewlyAdmitted bool Leases []ConcurrencyLease + Blockers []TaskAdmissionBlocker } type TaskAdmissionBatchOutcome struct { @@ -96,12 +110,15 @@ type TaskAdmissionBatchOutcome struct { } type TaskAdmissionMetricsSnapshot struct { - QueueDepth int - WaitingSync int - WaitingAsync int - OldestWaitSeconds float64 - ExpiredWaiterBacklog int - ExpiredDeadlineBacklog int + QueueDepth int + WaitingSync int + WaitingAsync int + DeferredAsync int + PendingReselect int + OldestWaitSeconds float64 + OldestDispatchableWaitSeconds float64 + ExpiredWaiterBacklog int + ExpiredDeadlineBacklog int } type TaskAdmissionReapResult struct { @@ -347,7 +364,7 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { } return taskAdmissionTxOutcome{Result: result}, nil } - if bindingChanged { + if bindingChanged || activeLeases == 0 { targetStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes) if stateErr != nil { return taskAdmissionTxOutcome{}, stateErr @@ -392,22 +409,31 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { } } - head, err := isTaskAdmissionHeadTx(ctx, tx, admission, scopes) + head, headBlockers, err := isTaskAdmissionHeadTx(ctx, tx, admission, scopes) if err != nil { return taskAdmissionTxOutcome{}, err } if !head { return taskAdmissionTxOutcome{ - Result: TaskAdmissionResult{Admission: admission}, + Result: TaskAdmissionResult{Admission: admission, Blockers: headBlockers}, }, nil } + blockers := make([]TaskAdmissionBlocker, 0, len(scopeStates)) for _, state := range scopeStates { if state.Saturated { - return taskAdmissionTxOutcome{ - Result: TaskAdmissionResult{Admission: admission}, - }, nil + blockers = append(blockers, TaskAdmissionBlocker{ + Reason: "saturated", + ScopeType: state.Scope.ScopeType, + ScopeKey: state.Scope.ScopeKey, + RetryAt: state.NextLeaseExpiration, + }) } } + if len(blockers) > 0 { + return taskAdmissionTxOutcome{ + Result: TaskAdmissionResult{Admission: admission, Blockers: blockers}, + }, nil + } leases := make([]ConcurrencyLease, 0, len(scopes)) for _, scope := range scopes { @@ -498,9 +524,6 @@ func (s *Store) TryTaskAdmissionBatchWithAdmittedHook( if admissionErr != nil { return outcomes, admissionErr } - if !result.Admitted { - break - } } return outcomes, nil } @@ -585,9 +608,6 @@ func (s *Store) tryTaskAdmissionAtomicBatchOnce( if outcome.NotifyTaskID != "" { notify = true } - if !outcome.Result.Admitted { - break - } } if err := tx.Commit(ctx); err != nil { return outcomes, err @@ -713,6 +733,11 @@ SET platform_id = $2::uuid, wait_deadline_at = LEAST(wait_deadline_at, $6), scope_snapshot = $7::jsonb, reselect_requested_at = NULL, + dispatch_next_at = now(), + dispatch_failure_count = 0, + dispatch_last_outcome = 'candidate_migrated', + dispatch_last_error_code = NULL, + dispatch_last_error_at = NULL, updated_at = now() WHERE task_id = $1::uuid AND status = 'waiting' @@ -720,7 +745,9 @@ RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at, - scope_snapshot, reselect_requested_at`, + scope_snapshot, reselect_requested_at, dispatch_next_at, + dispatch_failure_count, COALESCE(dispatch_last_outcome, ''), + COALESCE(dispatch_last_error_code, ''), dispatch_last_error_at`, input.TaskID, input.PlatformID, input.PlatformModelID, @@ -779,13 +806,20 @@ SET status = 'waiting', waiter_id = NULLIF($7, ''), waiter_lease_expires_at = $8, scope_snapshot = $9::jsonb, + dispatch_next_at = now(), + dispatch_failure_count = 0, + dispatch_last_outcome = NULL, + dispatch_last_error_code = NULL, + dispatch_last_error_at = NULL, updated_at = now() WHERE task_id = $1::uuid RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at, - scope_snapshot, reselect_requested_at`, + scope_snapshot, reselect_requested_at, dispatch_next_at, + dispatch_failure_count, COALESCE(dispatch_last_outcome, ''), + COALESCE(dispatch_last_error_code, ''), dispatch_last_error_at`, input.TaskID, input.PlatformID, input.PlatformModelID, @@ -915,6 +949,7 @@ LEFT JOIN river_job job ON job.id = task.river_job_id WHERE admission.mode = 'async' AND task.status = 'queued' AND task.next_run_at <= now() + AND admission.dispatch_next_at <= now() AND ( admission.status = 'waiting' OR ( @@ -964,13 +999,16 @@ SELECT admission.task_id::text, admission.platform_id::text, admission.platform_ COALESCE(admission.user_group_id::text, ''), admission.queue_key, admission.mode, admission.status, admission.priority, admission.enqueued_at, admission.wait_deadline_at, admission.admitted_at, COALESCE(admission.waiter_id, ''), admission.waiter_lease_expires_at, - admission.scope_snapshot, admission.reselect_requested_at + admission.scope_snapshot, admission.reselect_requested_at, admission.dispatch_next_at, + admission.dispatch_failure_count, COALESCE(admission.dispatch_last_outcome, ''), + COALESCE(admission.dispatch_last_error_code, ''), admission.dispatch_last_error_at 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.mode = 'async' AND task.status = 'queued' AND task.next_run_at <= now() + AND admission.dispatch_next_at <= now() AND ( admission.status = 'waiting' OR ( @@ -1006,27 +1044,105 @@ LIMIT $1`, limit) return admissions, rows.Err() } -// RequestWaitingTaskAdmissionReselect marks every queued task bound to a -// saturated platform model so the dispatcher can route the next batch to a -// different eligible candidate. The durable marker lets multiple dispatchers -// observe the same decision without assigning tasks to a specific Worker. -func (s *Store) RequestWaitingTaskAdmissionReselect(ctx context.Context, platformModelID string) (int64, error) { - result, err := s.pool.Exec(ctx, ` +// RequestTaskAdmissionReselect marks one waiting task for candidate +// reselection. The expected platform-model binding makes the operation a CAS: +// a concurrent dispatcher that already migrated or admitted the task wins. +func (s *Store) RequestTaskAdmissionReselect( + ctx context.Context, + taskID string, + expectedPlatformModelID string, +) (TaskAdmission, bool, error) { + row := s.pool.QueryRow(ctx, ` UPDATE gateway_task_admissions admission SET reselect_requested_at = now(), + dispatch_next_at = now(), + dispatch_last_outcome = 'reselect', + dispatch_last_error_code = NULL, + dispatch_last_error_at = NULL, updated_at = now() FROM gateway_tasks task -WHERE admission.task_id = task.id - AND admission.platform_model_id = $1::uuid +WHERE admission.task_id = $1::uuid + AND admission.task_id = task.id + AND admission.platform_model_id = $2::uuid AND admission.mode = 'async' AND admission.status = 'waiting' AND task.status = 'queued' AND task.next_run_at <= now() - AND admission.reselect_requested_at IS NULL`, platformModelID) - if err != nil { - return 0, err + AND admission.reselect_requested_at IS NULL +RETURNING admission.task_id::text, admission.platform_id::text, admission.platform_model_id::text, + COALESCE(admission.user_group_id::text, ''), admission.queue_key, admission.mode, + admission.status, admission.priority, admission.enqueued_at, admission.wait_deadline_at, + admission.admitted_at, COALESCE(admission.waiter_id, ''), admission.waiter_lease_expires_at, + admission.scope_snapshot, admission.reselect_requested_at, admission.dispatch_next_at, + admission.dispatch_failure_count, COALESCE(admission.dispatch_last_outcome, ''), + COALESCE(admission.dispatch_last_error_code, ''), admission.dispatch_last_error_at`, + taskID, + expectedPlatformModelID, + ) + admission, err := scanTaskAdmission(row) + if errors.Is(err, pgx.ErrNoRows) { + return TaskAdmission{}, false, nil } - return result.RowsAffected(), nil + if err != nil { + return TaskAdmission{}, false, err + } + s.notifyTaskAdmissionBestEffort(ctx, taskID) + return admission, true, nil +} + +// DeferWaitingTaskAdmission quarantines one task-specific dispatcher failure +// until retryAt. Deferred tasks are omitted from FIFO head selection so a +// corrupt or temporarily unroutable row cannot strand healthy followers. +func (s *Store) DeferWaitingTaskAdmission( + ctx context.Context, + taskID string, + expectedPlatformModelID string, + retryAt time.Time, + outcome string, + errorCode string, +) (TaskAdmission, bool, error) { + if retryAt.Before(time.Now().Add(time.Second)) { + retryAt = time.Now().Add(time.Second) + } + row := s.pool.QueryRow(ctx, ` +UPDATE gateway_task_admissions admission +SET dispatch_next_at = $3, + dispatch_failure_count = dispatch_failure_count + 1, + dispatch_last_outcome = NULLIF($4, ''), + dispatch_last_error_code = NULLIF($5, ''), + dispatch_last_error_at = now(), + updated_at = now() +FROM gateway_tasks task +WHERE admission.task_id = $1::uuid + AND admission.task_id = task.id + AND (NULLIF($2, '') IS NULL OR admission.platform_model_id = NULLIF($2, '')::uuid) + AND admission.mode = 'async' + AND admission.status = 'waiting' + AND task.status = 'queued' + AND task.next_run_at <= now() + AND admission.dispatch_next_at <= now() +RETURNING admission.task_id::text, admission.platform_id::text, admission.platform_model_id::text, + COALESCE(admission.user_group_id::text, ''), admission.queue_key, admission.mode, + admission.status, admission.priority, admission.enqueued_at, admission.wait_deadline_at, + admission.admitted_at, COALESCE(admission.waiter_id, ''), admission.waiter_lease_expires_at, + admission.scope_snapshot, admission.reselect_requested_at, admission.dispatch_next_at, + admission.dispatch_failure_count, COALESCE(admission.dispatch_last_outcome, ''), + COALESCE(admission.dispatch_last_error_code, ''), admission.dispatch_last_error_at`, + taskID, + strings.TrimSpace(expectedPlatformModelID), + retryAt, + strings.TrimSpace(outcome), + strings.TrimSpace(errorCode), + ) + admission, err := scanTaskAdmission(row) + if errors.Is(err, pgx.ErrNoRows) { + return TaskAdmission{}, false, nil + } + if err != nil { + return TaskAdmission{}, false, err + } + s.notifyTaskAdmissionBestEffort(ctx, taskID) + return admission, true, nil } // ListWaitingTaskAdmissionIDs returns one FIFO leader for every independent @@ -1094,7 +1210,16 @@ func (s *Store) TaskAdmissionMetrics(ctx context.Context) (TaskAdmissionMetricsS SELECT COUNT(*) FILTER (WHERE status = 'waiting')::int, COUNT(*) FILTER (WHERE status = 'waiting' AND mode = 'sync')::int, COUNT(*) FILTER (WHERE status = 'waiting' AND mode = 'async')::int, + COUNT(*) FILTER ( + WHERE status = 'waiting' AND mode = 'async' AND dispatch_next_at > now() + )::int, + COUNT(*) FILTER ( + WHERE status = 'waiting' AND mode = 'async' AND reselect_requested_at IS NOT NULL + )::int, COALESCE(EXTRACT(EPOCH FROM now() - MIN(enqueued_at) FILTER (WHERE status = 'waiting')), 0)::float8, + COALESCE(EXTRACT(EPOCH FROM now() - MIN(enqueued_at) FILTER ( + WHERE status = 'waiting' AND mode = 'async' AND dispatch_next_at <= now() + )), 0)::float8, COUNT(*) FILTER ( WHERE status = 'waiting' AND mode = 'sync' AND waiter_lease_expires_at <= now() )::int, @@ -1105,17 +1230,28 @@ FROM gateway_task_admissions`).Scan( &snapshot.QueueDepth, &snapshot.WaitingSync, &snapshot.WaitingAsync, + &snapshot.DeferredAsync, + &snapshot.PendingReselect, &snapshot.OldestWaitSeconds, + &snapshot.OldestDispatchableWaitSeconds, &snapshot.ExpiredWaiterBacklog, &snapshot.ExpiredDeadlineBacklog, ) return snapshot, err } -func (s *Store) ReapExpiredTaskAdmissions(ctx context.Context, limit int) (TaskAdmissionReapResult, error) { +func (s *Store) ReapExpiredTaskAdmissions( + ctx context.Context, + limit int, + defaultCallbackURLs ...string, +) (TaskAdmissionReapResult, error) { if limit <= 0 || limit > 1000 { limit = 500 } + defaultCallbackURL := "" + if len(defaultCallbackURLs) > 0 { + defaultCallbackURL = strings.TrimSpace(defaultCallbackURLs[0]) + } tx, err := s.pool.Begin(ctx) if err != nil { return TaskAdmissionReapResult{}, err @@ -1155,6 +1291,7 @@ LIMIT $1`, limit) return TaskAdmissionReapResult{}, err } var deadlineExpired bool + var dispatchFailureCount int err := tx.QueryRow(ctx, ` DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid @@ -1163,7 +1300,7 @@ WHERE task_id = $1::uuid wait_deadline_at <= now() OR (mode = 'sync' AND waiter_lease_expires_at <= now()) ) -RETURNING wait_deadline_at <= now()`, taskID).Scan(&deadlineExpired) +RETURNING wait_deadline_at <= now(), dispatch_failure_count`, taskID).Scan(&deadlineExpired, &dispatchFailureCount) if errors.Is(err, pgx.ErrNoRows) { continue } @@ -1177,22 +1314,112 @@ RETURNING wait_deadline_at <= now()`, taskID).Scan(&deadlineExpired) code = "queue_timeout" message = ErrQueueTimeout.Error() status = "failed" + if dispatchFailureCount > 0 { + code = "admission_dispatch_failed" + message = "task admission dispatcher could not prepare the task before its queue deadline" + } result.ExpiredDeadlines++ } else { result.ExpiredWaiters++ } - if _, err := tx.Exec(ctx, ` + publicErrorJSON := encodePublicErrorSnapshot(code, message, 0, true, "", taskID) + tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = $2, - error = $3, + error = NULL, error_code = $4, error_message = $3, + public_error = $5::jsonb, + billing_status = CASE + WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required' + WHEN reservation_amount > 0 THEN 'pending' + ELSE 'released' + END, + billing_updated_at = now(), + locked_by = NULL, + locked_at = NULL, + heartbeat_at = NULL, + execution_token = NULL, + execution_lease_expires_at = NULL, finished_at = now(), updated_at = now() WHERE id = $1::uuid - AND status = 'queued'`, taskID, status, message, code); err != nil { + AND status = 'queued'`, taskID, status, message, code, string(publicErrorJSON)) + if err != nil { return TaskAdmissionReapResult{}, err } + if tag.RowsAffected() == 0 { + continue + } + if _, err := tx.Exec(ctx, ` +UPDATE gateway_concurrency_leases +SET released_at = statement_timestamp() +WHERE task_id = $1::uuid + AND released_at IS NULL`, taskID); err != nil { + return TaskAdmissionReapResult{}, err + } + if _, err := releaseRuntimeRateReservationsTx(ctx, tx, []string{taskID}); err != nil { + return TaskAdmissionReapResult{}, err + } + payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": code}) + if _, err := tx.Exec(ctx, ` +INSERT INTO settlement_outbox ( + task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at +) +SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency, + pricing_snapshot, $2::jsonb, 'pending', now() +FROM gateway_tasks +WHERE id = $1::uuid + AND run_mode IN ('production', 'acceptance', 'acceptance_canary') + AND gateway_user_id IS NOT NULL + AND reservation_amount > 0 +ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil { + return TaskAdmissionReapResult{}, err + } + eventType := "task.cancelled" + if status == "failed" { + eventType = "task.failed" + } + var eventID string + var eventSeq int64 + if err := tx.QueryRow(ctx, ` +WITH next_seq AS ( + SELECT COALESCE(MAX(seq), 0) + 1 AS seq + FROM gateway_task_events + WHERE task_id = $1::uuid +) +INSERT INTO gateway_task_events ( + task_id, seq, event_type, status, phase, progress, message, payload, simulated +) +SELECT $1::uuid, next_seq.seq, $2, $3, 'admission', 0, $4, '{}'::jsonb, + task.run_mode = 'simulation' +FROM next_seq +JOIN gateway_tasks task ON task.id = $1::uuid +RETURNING id::text, seq`, taskID, eventType, status, message).Scan(&eventID, &eventSeq); err != nil { + return TaskAdmissionReapResult{}, err + } + var acceptanceCallbackURL string + if err := tx.QueryRow(ctx, ` +SELECT COALESCE(( + SELECT run.callback_url + FROM gateway_tasks task + JOIN gateway_acceptance_runs run ON run.id = task.acceptance_run_id + WHERE task.id = $1::uuid +), '')`, taskID).Scan(&acceptanceCallbackURL); err != nil { + return TaskAdmissionReapResult{}, err + } + callbackURL := defaultCallbackURL + if strings.TrimSpace(acceptanceCallbackURL) != "" { + callbackURL = strings.TrimSpace(acceptanceCallbackURL) + } + if callbackURL != "" { + if _, err := tx.Exec(ctx, ` +INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload) +VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb) +ON CONFLICT (task_id, seq, callback_url) DO NOTHING`, taskID, eventID, eventSeq, callbackURL); err != nil { + return TaskAdmissionReapResult{}, err + } + } reaped++ } if err := tx.Commit(ctx); err != nil { @@ -1448,7 +1675,9 @@ RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at, - scope_snapshot, reselect_requested_at`, + scope_snapshot, reselect_requested_at, dispatch_next_at, + dispatch_failure_count, COALESCE(dispatch_last_outcome, ''), + COALESCE(dispatch_last_error_code, ''), dispatch_last_error_at`, input.TaskID, input.PlatformID, input.PlatformModelID, input.UserGroupID, input.QueueKey, input.Mode, input.Priority, deadline, waiterID, waiterLease, scopeSnapshot) return scanTaskAdmission(row) @@ -1467,7 +1696,9 @@ RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at, - scope_snapshot, reselect_requested_at`, + scope_snapshot, reselect_requested_at, dispatch_next_at, + dispatch_failure_count, COALESCE(dispatch_last_outcome, ''), + COALESCE(dispatch_last_error_code, ''), dispatch_last_error_at`, taskID, waiterID, admissionWaiterLeaseTTL.String()) return scanTaskAdmission(row) } @@ -1478,13 +1709,19 @@ UPDATE gateway_task_admissions SET status = 'admitted', admitted_at = now(), waiter_lease_expires_at = NULL, + dispatch_failure_count = 0, + dispatch_last_outcome = 'admitted', + dispatch_last_error_code = NULL, + dispatch_last_error_at = NULL, updated_at = now() WHERE task_id = $1::uuid AND status = 'waiting' RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at, - scope_snapshot, reselect_requested_at`, + scope_snapshot, reselect_requested_at, dispatch_next_at, + dispatch_failure_count, COALESCE(dispatch_last_outcome, ''), + COALESCE(dispatch_last_error_code, ''), dispatch_last_error_at`, taskID) return scanTaskAdmission(row) } @@ -1503,7 +1740,9 @@ SELECT task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at, - scope_snapshot, reselect_requested_at + scope_snapshot, reselect_requested_at, dispatch_next_at, + dispatch_failure_count, COALESCE(dispatch_last_outcome, ''), + COALESCE(dispatch_last_error_code, ''), dispatch_last_error_at FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID) admission, err := scanTaskAdmission(row) @@ -1518,6 +1757,7 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) { var admittedAt sql.NullTime var waiterLeaseExpires sql.NullTime var reselectRequestedAt sql.NullTime + var dispatchLastErrorAt sql.NullTime var scopeSnapshot []byte err := row.Scan( &admission.TaskID, @@ -1535,6 +1775,11 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) { &waiterLeaseExpires, &scopeSnapshot, &reselectRequestedAt, + &admission.DispatchNextAt, + &admission.DispatchFailureCount, + &admission.DispatchLastOutcome, + &admission.DispatchLastErrorCode, + &dispatchLastErrorAt, ) if admittedAt.Valid { admission.AdmittedAt = admittedAt.Time @@ -1545,15 +1790,18 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) { if reselectRequestedAt.Valid { admission.ReselectRequestedAt = reselectRequestedAt.Time } + if dispatchLastErrorAt.Valid { + admission.DispatchLastErrorAt = dispatchLastErrorAt.Time + } if err == nil && len(scopeSnapshot) > 0 { if unmarshalErr := json.Unmarshal(scopeSnapshot, &admission.Scopes); unmarshalErr != nil { - return TaskAdmission{}, fmt.Errorf("unmarshal task admission scopes: %w", unmarshalErr) + admission.SnapshotInvalid = true } } return admission, err } -func isTaskAdmissionHeadTx(ctx context.Context, tx pgx.Tx, admission TaskAdmission, scopes []AdmissionScope) (bool, error) { +func isTaskAdmissionHeadTx(ctx context.Context, tx pgx.Tx, admission TaskAdmission, scopes []AdmissionScope) (bool, []TaskAdmissionBlocker, error) { for _, scope := range scopes { var head string switch scope.ScopeType { @@ -1561,10 +1809,12 @@ func isTaskAdmissionHeadTx(ctx context.Context, tx pgx.Tx, admission TaskAdmissi if err := tx.QueryRow(ctx, ` SELECT task_id::text FROM gateway_task_admissions -WHERE platform_model_id = $1::uuid AND status = 'waiting' +WHERE platform_model_id = $1::uuid + AND status = 'waiting' + AND (mode <> 'async' OR dispatch_next_at <= statement_timestamp()) ORDER BY priority ASC, enqueued_at ASC, task_id ASC LIMIT 1`, admission.PlatformModelID).Scan(&head); err != nil { - return false, err + return false, nil, err } case "user_group": if admission.UserGroupID == "" { @@ -1573,19 +1823,25 @@ LIMIT 1`, admission.PlatformModelID).Scan(&head); err != nil { if err := tx.QueryRow(ctx, ` SELECT task_id::text FROM gateway_task_admissions -WHERE user_group_id = $1::uuid AND status = 'waiting' +WHERE user_group_id = $1::uuid + AND status = 'waiting' + AND (mode <> 'async' OR dispatch_next_at <= statement_timestamp()) ORDER BY priority ASC, enqueued_at ASC, task_id ASC LIMIT 1`, admission.UserGroupID).Scan(&head); err != nil { - return false, err + return false, nil, err } default: continue } if head != admission.TaskID { - return false, nil + return false, []TaskAdmissionBlocker{{ + Reason: "not_queue_head", + ScopeType: scope.ScopeType, + ScopeKey: scope.ScopeKey, + }}, nil } } - return true, nil + return true, nil, nil } func expireTaskAdmissionTx(ctx context.Context, tx pgx.Tx, taskID string) error { diff --git a/apps/api/internal/store/admission_queue_integration_test.go b/apps/api/internal/store/admission_queue_integration_test.go index ce80ae3..761a080 100644 --- a/apps/api/internal/store/admission_queue_integration_test.go +++ b/apps/api/internal/store/admission_queue_integration_test.go @@ -131,6 +131,9 @@ RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID); if err != nil || result.Admitted { t.Fatalf("second task should wait: result=%+v err=%v", result, err) } + if len(result.Blockers) != 1 || result.Blockers[0].Reason != "saturated" || result.Blockers[0].ScopeType != "platform_model" { + t.Fatalf("second task blockers=%+v, want saturated platform_model", result.Blockers) + } higherPriorityAsync := createTask(true) result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, "")) if err != nil || result.Admitted { @@ -176,6 +179,9 @@ WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil { if result.Admitted { t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter") } + if len(result.Blockers) != 1 || result.Blockers[0].Reason != "not_queue_head" || result.Blockers[0].ScopeType != "platform_model" { + t.Fatalf("lower-priority blockers=%+v, want platform FIFO head blocker", result.Blockers) + } result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, "")) if err != nil || !result.Admitted { t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err) @@ -451,12 +457,23 @@ WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID) if listedSnapshot == nil || len(listedSnapshot.Scopes) != len(queuedAdmission.Scopes) { t.Fatalf("listed admission snapshot=%+v, want %d scopes", listedSnapshot, len(queuedAdmission.Scopes)) } - markedForReselect, err := first.RequestWaitingTaskAdmissionReselect(ctx, platformModelID) + _, markedForReselect, err := first.RequestTaskAdmissionReselect(ctx, queuedAtomicTask.ID, platformModelID) if err != nil { t.Fatalf("request waiting admission reselection: %v", err) } - if markedForReselect < 1 { - t.Fatalf("marked admissions=%d, want at least the queued task", markedForReselect) + if !markedForReselect { + t.Fatal("queued task was not marked for reselection") + } + var platformReselectCount int + if err := first.pool.QueryRow(ctx, ` +SELECT count(*) +FROM gateway_task_admissions +WHERE platform_model_id = $1::uuid + AND reselect_requested_at IS NOT NULL`, platformModelID).Scan(&platformReselectCount); err != nil { + t.Fatalf("count task-scoped reselections: %v", err) + } + if platformReselectCount != 1 { + t.Fatalf("platform reselection fanout=%d, want exactly one task", platformReselectCount) } reselectAdmission, err := first.GetTaskAdmission(ctx, queuedAtomicTask.ID) if err != nil { @@ -1070,6 +1087,276 @@ WHERE id = $1::uuid`, ambiguousTask.ID).Scan(&ambiguousStatus); err != nil { t.Fatalf("ambiguous submission task status = %s, want running", ambiguousStatus) } + globalHolder := createTask(true) + globalWaiter := createTask(true) + globalScope := AdmissionScope{ + ScopeType: "worker_capacity", + ScopeKey: "global-blocker-" + suffix, + ScopeName: "global blocker", + ConcurrentLimit: 1, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 10, + MaxWaitSeconds: 600, + } + globalHolderInput := inputFor(globalHolder, 100, "") + globalHolderInput.Scopes = []AdmissionScope{globalScope} + globalWaiterInput := inputFor(globalWaiter, 100, "") + globalWaiterInput.Scopes = []AdmissionScope{globalScope} + if result, err := first.TryTaskAdmission(ctx, globalHolderInput); err != nil || !result.Admitted { + t.Fatalf("admit global capacity holder: result=%+v err=%v", result, err) + } + if _, err := first.QueueTaskAdmissionWithHook(ctx, globalWaiterInput, nil); err != nil { + t.Fatalf("queue global capacity waiter: %v", err) + } + globalWaitResult, err := second.TryTaskAdmission(ctx, globalWaiterInput) + if err != nil || globalWaitResult.Admitted || len(globalWaitResult.Blockers) != 1 || + globalWaitResult.Blockers[0].Reason != "saturated" || + globalWaitResult.Blockers[0].ScopeType != "worker_capacity" { + t.Fatalf("global capacity waiter result=%+v err=%v", globalWaitResult, err) + } + if !globalWaitResult.Admission.ReselectRequestedAt.IsZero() { + t.Fatal("global capacity wait unexpectedly requested candidate reselection") + } + if err := first.DeleteTaskAdmission(ctx, globalHolder.ID); err != nil { + t.Fatalf("release global capacity holder: %v", err) + } + if err := first.DeleteTaskAdmission(ctx, globalWaiter.ID); err != nil { + t.Fatalf("delete global capacity waiter: %v", err) + } + + orphanedAdmittedTask := createTask(true) + orphanedAdmittedInput := inputFor(orphanedAdmittedTask, 100, "") + orphanedAdmittedInput.Scopes = []AdmissionScope{{ + ScopeType: "worker_capacity", + ScopeKey: "orphaned-admitted-" + suffix, + ConcurrentLimit: 1, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 10, + MaxWaitSeconds: 600, + }} + if result, err := first.TryTaskAdmission(ctx, orphanedAdmittedInput); err != nil || !result.Admitted { + t.Fatalf("admit orphan candidate: result=%+v err=%v", result, err) + } + if _, err := first.pool.Exec(ctx, ` +UPDATE gateway_concurrency_leases +SET released_at = now() +WHERE task_id = $1::uuid + AND released_at IS NULL`, orphanedAdmittedTask.ID); err != nil { + t.Fatalf("orphan admitted task leases: %v", err) + } + if result, err := second.TryTaskAdmission(ctx, orphanedAdmittedInput); err != nil || !result.Admitted || !result.NewlyAdmitted { + t.Fatalf("recover admitted row without lease: result=%+v err=%v", result, err) + } + if err := first.DeleteTaskAdmission(ctx, orphanedAdmittedTask.ID); err != nil { + t.Fatalf("delete recovered admitted row: %v", err) + } + + resilienceTask := createTask(true) + resilienceInput := inputFor(resilienceTask, 100, "") + resilienceInput.Scopes = []AdmissionScope{{ + ScopeType: "worker_capacity", + ScopeKey: "resilience-" + suffix, + ScopeName: "resilience capacity", + ConcurrentLimit: 5, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 10, + MaxWaitSeconds: 600, + }} + if _, err := first.QueueTaskAdmissionWithHook(ctx, resilienceInput, nil); err != nil { + t.Fatalf("queue resilience task: %v", err) + } + if _, err := first.pool.Exec(ctx, ` +UPDATE gateway_task_admissions +SET scope_snapshot = '[{"ScopeType":123}]'::jsonb +WHERE task_id = $1::uuid`, resilienceTask.ID); err != nil { + t.Fatalf("corrupt resilience scope snapshot: %v", err) + } + listedAdmissions, err := first.ListWaitingAsyncAdmissions(ctx, 1000) + if err != nil { + t.Fatalf("list admissions with one corrupt snapshot: %v", err) + } + corruptFound := false + for _, listed := range listedAdmissions { + if listed.TaskID == resilienceTask.ID { + corruptFound = listed.SnapshotInvalid + break + } + } + if !corruptFound { + t.Fatal("corrupt admission snapshot was not isolated for rebuilding") + } + repairedAdmission, err := first.RebindWaitingTaskAdmission(ctx, resilienceInput) + if err != nil || repairedAdmission.SnapshotInvalid || len(repairedAdmission.Scopes) != 1 { + t.Fatalf("repair corrupt admission snapshot=%+v err=%v", repairedAdmission, err) + } + deferred, changed, err := first.DeferWaitingTaskAdmission( + ctx, + resilienceTask.ID, + platformModelID, + time.Now().Add(time.Minute), + "deferred", + "snapshot_invalid", + ) + if err != nil || !changed || deferred.DispatchFailureCount != 1 || deferred.DispatchLastErrorCode != "snapshot_invalid" { + t.Fatalf("defer corrupt admission=%+v changed=%v err=%v", deferred, changed, err) + } + listedAdmissions, err = first.ListWaitingAsyncAdmissions(ctx, 1000) + if err != nil { + t.Fatalf("list admissions after defer: %v", err) + } + for _, listed := range listedAdmissions { + if listed.TaskID == resilienceTask.ID { + t.Fatal("deferred admission remained dispatchable") + } + } + if _, err := first.pool.Exec(ctx, ` +UPDATE gateway_task_admissions +SET dispatch_next_at = now(), scope_snapshot = '[]'::jsonb +WHERE task_id = $1::uuid`, resilienceTask.ID); err != nil { + t.Fatalf("make resilience admission dispatchable: %v", err) + } + reselected, changed, err := first.RequestTaskAdmissionReselect(ctx, resilienceTask.ID, platformModelID) + if err != nil || !changed || reselected.ReselectRequestedAt.IsZero() { + t.Fatalf("task-scoped reselection=%+v changed=%v err=%v", reselected, changed, err) + } + if _, changed, err := second.RequestTaskAdmissionReselect(ctx, resilienceTask.ID, platformModelID); err != nil || changed { + t.Fatalf("duplicate task-scoped reselection changed=%v err=%v", changed, err) + } + if _, err := first.FailQueuedTaskWithCallback( + ctx, + resilienceTask.ID, + "no_model_candidate", + "no enabled platform model matches request", + "https://callback.invalid/admission-resilience", + false, + ); err != nil { + t.Fatalf("terminally fail resilience task: %v", err) + } + if _, err := second.FailQueuedTaskWithCallback( + ctx, + resilienceTask.ID, + "no_model_candidate", + "duplicate terminalization", + "https://callback.invalid/admission-resilience", + false, + ); !errors.Is(err, ErrTaskExecutionFinished) { + t.Fatalf("duplicate terminalization error=%v, want task finished", err) + } + var resilienceStatus, resilienceCode string + var resilienceAdmissions, resilienceLeases, resilienceFailedEvents, resilienceCallbacks int + if err := first.pool.QueryRow(ctx, ` +SELECT task.status, COALESCE(task.error_code, ''), + (SELECT count(*) FROM gateway_task_admissions admission WHERE admission.task_id = task.id), + (SELECT count(*) FROM gateway_concurrency_leases lease WHERE lease.task_id = task.id AND lease.released_at IS NULL), + (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.failed'), + (SELECT count(*) FROM gateway_task_callback_outbox callback WHERE callback.task_id = task.id) +FROM gateway_tasks task +WHERE task.id = $1::uuid`, resilienceTask.ID).Scan( + &resilienceStatus, + &resilienceCode, + &resilienceAdmissions, + &resilienceLeases, + &resilienceFailedEvents, + &resilienceCallbacks, + ); err != nil { + t.Fatalf("read resilience terminal state: %v", err) + } + if resilienceStatus != "failed" || resilienceCode != "no_model_candidate" || + resilienceAdmissions != 0 || resilienceLeases != 0 || + resilienceFailedEvents != 1 || resilienceCallbacks != 1 { + t.Fatalf( + "resilience terminal state=%s/%s admissions=%d leases=%d events=%d callbacks=%d", + resilienceStatus, + resilienceCode, + resilienceAdmissions, + resilienceLeases, + resilienceFailedEvents, + resilienceCallbacks, + ) + } + retryExpiredTask := createTask(true) + retryExpiredInput := resilienceInput + retryExpiredInput.TaskID = retryExpiredTask.ID + if _, err := first.QueueTaskAdmissionWithHook(ctx, retryExpiredInput, nil); err != nil { + t.Fatalf("queue retry-expired task: %v", err) + } + if _, err := first.pool.Exec(ctx, ` +UPDATE gateway_task_admissions +SET enqueued_at = now() - interval '2 seconds', + wait_deadline_at = now() - interval '1 second', + dispatch_failure_count = 3, + dispatch_last_outcome = 'deferred', + dispatch_last_error_code = 'client_error', + dispatch_last_error_at = now() - interval '1 minute' +WHERE task_id = $1::uuid`, retryExpiredTask.ID); err != nil { + t.Fatalf("expire deferred task: %v", err) + } + reaped, err = first.ReapExpiredTaskAdmissions(ctx, 10, "https://callback.invalid/admission-timeout") + if err != nil || reaped.ExpiredDeadlines < 1 { + t.Fatalf("reap retry-expired task=%+v err=%v", reaped, err) + } + var retryExpiredCode string + var retryExpiredEvents, retryExpiredCallbacks int + if err := first.pool.QueryRow(ctx, ` +SELECT COALESCE(task.error_code, ''), + (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.failed'), + (SELECT count(*) FROM gateway_task_callback_outbox callback WHERE callback.task_id = task.id) +FROM gateway_tasks task +WHERE task.id = $1::uuid`, retryExpiredTask.ID).Scan( + &retryExpiredCode, + &retryExpiredEvents, + &retryExpiredCallbacks, + ); err != nil { + t.Fatalf("read retry-expired task: %v", err) + } + if retryExpiredCode != "admission_dispatch_failed" || retryExpiredEvents != 1 || retryExpiredCallbacks != 1 { + t.Fatalf( + "retry-expired task code=%s events=%d callbacks=%d, want admission_dispatch_failed/1/1", + retryExpiredCode, + retryExpiredEvents, + retryExpiredCallbacks, + ) + } + + disconnectedTask := createTask(true) + disconnectedInput := resilienceInput + disconnectedInput.TaskID = disconnectedTask.ID + if _, err := first.QueueTaskAdmissionWithHook(ctx, disconnectedInput, nil); err != nil { + t.Fatalf("queue disconnected task: %v", err) + } + if _, changed, err := first.CancelQueuedTaskWithCallback( + ctx, + disconnectedTask.ID, + "client disconnected before upstream submission", + "https://callback.invalid/admission-disconnect", + ); err != nil || !changed { + t.Fatalf("cancel disconnected task changed=%v err=%v", changed, err) + } + var disconnectedEvents, disconnectedCallbacks, disconnectedAdmissions int + if err := first.pool.QueryRow(ctx, ` +SELECT (SELECT count(*) FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.cancelled'), + (SELECT count(*) FROM gateway_task_callback_outbox callback WHERE callback.task_id = task.id), + (SELECT count(*) FROM gateway_task_admissions admission WHERE admission.task_id = task.id) +FROM gateway_tasks task +WHERE task.id = $1::uuid`, disconnectedTask.ID).Scan( + &disconnectedEvents, + &disconnectedCallbacks, + &disconnectedAdmissions, + ); err != nil { + t.Fatalf("read disconnected task cleanup: %v", err) + } + if disconnectedEvents != 1 || disconnectedCallbacks != 1 || disconnectedAdmissions != 0 { + t.Fatalf( + "disconnected task events=%d callbacks=%d admissions=%d, want 1/1/0", + disconnectedEvents, + disconnectedCallbacks, + disconnectedAdmissions, + ) + } + terminalResidue := createTask(true) result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, "")) if err != nil || !result.Admitted { @@ -1247,6 +1534,150 @@ WHERE task_id = ANY($1::uuid[])`, taskIDs); err != nil { if !leaderSet[tasks[0].ID] || leaderSet[tasks[2].ID] { t.Fatalf("grouped candidate leaders = %v, want only group head %s", leaders, tasks[0].ID) } + if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = ANY($1::uuid[])`, taskIDs); err != nil { + t.Fatalf("clear grouped admission leaders: %v", err) + } + + createAsyncTask := func(modelIndex int) GatewayTask { + t.Helper() + task, createErr := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generate", + Model: fmt.Sprintf("leader-model-%d-%s", modelIndex, suffix), + RunMode: "production", + Async: true, + Request: map[string]any{"prompt": "async admission independence test"}, + }, &auth.User{ID: "leader-test-user-" + suffix, Source: "gateway"}) + if createErr != nil { + t.Fatalf("create async task: %v", createErr) + } + taskIDs = append(taskIDs, task.ID) + return task + } + platformScope := func(modelID string) AdmissionScope { + return AdmissionScope{ + ScopeType: "platform_model", + ScopeKey: modelID, + ConcurrentLimit: 1, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 10, + MaxWaitSeconds: 600, + } + } + groupScope := AdmissionScope{ + ScopeType: "user_group", + ScopeKey: group.ID, + ConcurrentLimit: 1, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 10, + MaxWaitSeconds: 600, + } + holder := createAsyncTask(0) + blocked := createAsyncTask(0) + independent := createAsyncTask(1) + holderInput := TaskAdmissionInput{ + TaskID: holder.ID, PlatformID: platform.ID, PlatformModelID: modelIDs[0], + UserGroupID: group.ID, QueueKey: "leader-test:" + modelIDs[0], Mode: "async", Priority: 100, + Scopes: []AdmissionScope{platformScope(modelIDs[0]), groupScope}, + } + blockedInput := holderInput + blockedInput.TaskID = blocked.ID + independentInput := TaskAdmissionInput{ + TaskID: independent.ID, PlatformID: platform.ID, PlatformModelID: modelIDs[1], + QueueKey: "leader-test:" + modelIDs[1], Mode: "async", Priority: 100, + Scopes: []AdmissionScope{platformScope(modelIDs[1])}, + } + if result, err := db.TryTaskAdmission(ctx, holderInput); err != nil || !result.Admitted { + t.Fatalf("admit local blocker holder: result=%+v err=%v", result, err) + } + if _, err := db.QueueTaskAdmissionWithHook(ctx, blockedInput, nil); err != nil { + t.Fatalf("queue locally blocked task: %v", err) + } + if _, err := db.QueueTaskAdmissionWithHook(ctx, independentInput, nil); err != nil { + t.Fatalf("queue independent task: %v", err) + } + outcomes, err := db.TryTaskAdmissionAtomicBatchWithAdmittedHook( + ctx, + []TaskAdmissionInput{blockedInput, independentInput}, + nil, + ) + if err != nil || len(outcomes) != 2 { + t.Fatalf("independent atomic outcomes=%+v err=%v", outcomes, err) + } + if outcomes[0].Result.Admitted || len(outcomes[0].Result.Blockers) == 0 { + t.Fatalf("locally blocked outcome=%+v, want explicit blockers", outcomes[0]) + } + if !outcomes[1].Result.Admitted { + t.Fatalf("independent candidate was cross-blocked: %+v", outcomes[1]) + } + for _, taskID := range []string{holder.ID, blocked.ID, independent.ID} { + if err := db.DeleteTaskAdmission(ctx, taskID); err != nil { + t.Fatalf("clear independent atomic admission %s: %v", taskID, err) + } + } + + hookTasks := []GatewayTask{createAsyncTask(0), createAsyncTask(0), createAsyncTask(1)} + hookInputs := make([]TaskAdmissionInput, 0, len(hookTasks)) + for index, task := range hookTasks { + input := TaskAdmissionInput{ + TaskID: task.ID, PlatformID: platform.ID, PlatformModelID: modelIDs[index%len(modelIDs)], + QueueKey: "hook-isolation:" + task.ID, Mode: "async", Priority: 100, + Scopes: []AdmissionScope{{ + ScopeType: "worker_capacity", + ScopeKey: "hook-isolation-" + task.ID, + ConcurrentLimit: 1, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 10, + MaxWaitSeconds: 600, + }}, + } + if _, err := db.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil { + t.Fatalf("queue hook isolation task %d: %v", index, err) + } + hookInputs = append(hookInputs, input) + } + hookFailure := errors.New("synthetic per-task hook failure") + outcomes, err = db.TryTaskAdmissionAtomicBatchWithAdmittedHook( + ctx, + hookInputs, + func(_ pgx.Tx, input TaskAdmissionInput) error { + if input.TaskID == hookTasks[1].ID { + return hookFailure + } + return nil + }, + ) + if !errors.Is(err, hookFailure) || len(outcomes) != 2 || outcomes[1].TaskID != hookTasks[1].ID { + t.Fatalf("hook failure outcomes=%+v err=%v", outcomes, err) + } + var hookAdmitted, hookLeases int + if err := db.pool.QueryRow(ctx, ` +SELECT (SELECT count(*) FROM gateway_task_admissions WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'), + (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL)`, + []string{hookTasks[0].ID, hookTasks[1].ID, hookTasks[2].ID}, + ).Scan(&hookAdmitted, &hookLeases); err != nil { + t.Fatalf("read rolled back hook batch: %v", err) + } + if hookAdmitted != 0 || hookLeases != 0 { + t.Fatalf("failed hook batch left admitted=%d leases=%d", hookAdmitted, hookLeases) + } + if _, changed, err := db.DeferWaitingTaskAdmission( + ctx, + hookTasks[1].ID, + hookInputs[1].PlatformModelID, + time.Now().Add(time.Minute), + "deferred", + "synthetic_hook_failure", + ); err != nil || !changed { + t.Fatalf("isolate hook failure changed=%v err=%v", changed, err) + } + for _, index := range []int{0, 2} { + if result, err := db.TryTaskAdmission(ctx, hookInputs[index]); err != nil || !result.Admitted { + t.Fatalf("admit hook survivor %d result=%+v err=%v", index, result, err) + } + } } func TestWorkerCapacityAllocationAndFailover(t *testing.T) { diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index 345ef04..7e55ffc 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -545,17 +545,40 @@ WHERE id = $1::uuid } func (s *Store) FailQueuedTask(ctx context.Context, taskID string, code string, message string) (GatewayTask, error) { + return s.FailQueuedTaskWithCallback(ctx, taskID, code, message, "", false) +} + +// FailQueuedTaskWithCallback terminally fails a task before upstream +// submission and atomically releases every durable resource owned by its +// admission. The task advisory lock makes competing dispatchers idempotent. +func (s *Store) FailQueuedTaskWithCallback( + ctx context.Context, + taskID string, + code string, + message string, + callbackURL string, + simulated bool, +) (GatewayTask, error) { + code = strings.TrimSpace(code) + message = truncateUTF8Bytes(message, 2048) publicErrorJSON := encodePublicErrorSnapshot(code, message, 0, true, "", taskID) - tag, err := s.pool.Exec(ctx, ` + var task GatewayTask + err := s.beginTransaction(ctx, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil { + return err + } + var err error + task, err = scanGatewayTask(tx.QueryRow(ctx, ` UPDATE gateway_tasks SET status = 'failed', error = NULL, error_code = NULLIF($2, ''), error_message = NULLIF($3, ''), - public_error = $4::jsonb, + public_error = $4::jsonb, billing_status = CASE WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required' - ELSE 'released' + WHEN reservation_amount > 0 THEN 'pending' + ELSE 'released' END, billing_updated_at = now(), locked_by = NULL, @@ -566,14 +589,76 @@ SET status = 'failed', finished_at = now(), updated_at = now() WHERE id = $1::uuid - AND status = 'queued'`, taskID, strings.TrimSpace(code), truncateUTF8Bytes(message, 2048), string(publicErrorJSON)) + AND status = 'queued' + AND COALESCE(remote_task_id, '') = '' +RETURNING `+gatewayTaskColumns, taskID, code, message, string(publicErrorJSON))) + if errors.Is(err, pgx.ErrNoRows) { + return ErrTaskExecutionFinished + } + if err != nil { + return err + } + if _, err := tx.Exec(ctx, ` +UPDATE gateway_concurrency_leases +SET released_at = statement_timestamp() +WHERE task_id = $1::uuid + AND released_at IS NULL`, taskID); err != nil { + return err + } + if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { + return err + } + if _, err := releaseRuntimeRateReservationsTx(ctx, tx, []string{taskID}); err != nil { + return err + } + payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": code}) + if _, err := tx.Exec(ctx, ` +INSERT INTO settlement_outbox ( + task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at +) +SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency, + pricing_snapshot, $2::jsonb, 'pending', now() +FROM gateway_tasks +WHERE id = $1::uuid + AND run_mode IN ('production', 'acceptance', 'acceptance_canary') + AND gateway_user_id IS NOT NULL + AND reservation_amount > 0 +ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil { + return err + } + var eventID string + var eventSeq int64 + if err := tx.QueryRow(ctx, ` +WITH next_seq AS ( + SELECT COALESCE(MAX(seq), 0) + 1 AS seq + FROM gateway_task_events + WHERE task_id = $1::uuid +) +INSERT INTO gateway_task_events ( + task_id, seq, event_type, status, phase, progress, message, payload, simulated +) +SELECT $1::uuid, next_seq.seq, 'task.failed', 'failed', 'admission', 0, + NULLIF($2, ''), '{}'::jsonb, $3 +FROM next_seq +RETURNING id::text, seq`, taskID, message, simulated).Scan(&eventID, &eventSeq); err != nil { + return err + } + callbackURL = strings.TrimSpace(callbackURL) + if callbackURL != "" { + if _, err := tx.Exec(ctx, ` +INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload) +VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb) +ON CONFLICT (task_id, seq, callback_url) DO NOTHING`, taskID, eventID, eventSeq, callbackURL); err != nil { + return err + } + } + return nil + }) if err != nil { return GatewayTask{}, err } - if tag.RowsAffected() != 1 { - return GatewayTask{}, ErrTaskExecutionFinished - } - return s.GetTask(ctx, taskID) + s.notifyTaskAdmissionBestEffort(ctx, "*") + return task, nil } func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error { @@ -888,6 +973,25 @@ WHERE attempt.id = $1::uuid } func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message string) (GatewayTask, bool, error) { + return s.cancelQueuedTask(ctx, taskID, message, "", false) +} + +func (s *Store) CancelQueuedTaskWithCallback( + ctx context.Context, + taskID string, + message string, + callbackURL string, +) (GatewayTask, bool, error) { + return s.cancelQueuedTask(ctx, taskID, message, callbackURL, true) +} + +func (s *Store) cancelQueuedTask( + ctx context.Context, + taskID string, + message string, + callbackURL string, + emitTerminalEvent bool, +) (GatewayTask, bool, error) { message = strings.TrimSpace(message) if message == "" { message = "任务已取消" @@ -941,8 +1045,11 @@ WHERE task_id = $1::uuid if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { return err } + if _, err := releaseRuntimeRateReservationsTx(ctx, tx, []string{taskID}); err != nil { + return err + } payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "queued_cancelled"}) - _, err = tx.Exec(ctx, ` + if _, err = tx.Exec(ctx, ` INSERT INTO settlement_outbox ( task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at ) @@ -953,7 +1060,37 @@ WHERE id = $1::uuid AND run_mode IN ('production', 'acceptance', 'acceptance_canary') AND gateway_user_id IS NOT NULL AND reservation_amount > 0 -ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)) +ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil { + return err + } + if !emitTerminalEvent { + return nil + } + var eventID string + var eventSeq int64 + if err := tx.QueryRow(ctx, ` +WITH next_seq AS ( + SELECT COALESCE(MAX(seq), 0) + 1 AS seq + FROM gateway_task_events + WHERE task_id = $1::uuid +) +INSERT INTO gateway_task_events ( + task_id, seq, event_type, status, phase, progress, message, payload, simulated +) +SELECT $1::uuid, next_seq.seq, 'task.cancelled', 'cancelled', 'admission', 0, + NULLIF($2, ''), '{}'::jsonb, $3 +FROM next_seq +RETURNING id::text, seq`, taskID, message, task.RunMode == "simulation").Scan(&eventID, &eventSeq); err != nil { + return err + } + callbackURL = strings.TrimSpace(callbackURL) + if callbackURL == "" { + return nil + } + _, err = tx.Exec(ctx, ` +INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload) +VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb) +ON CONFLICT (task_id, seq, callback_url) DO NOTHING`, taskID, eventID, eventSeq, callbackURL) return err }) if err != nil { diff --git a/apps/api/internal/workerload/controller.go b/apps/api/internal/workerload/controller.go index 9d85f56..98ffe3e 100644 --- a/apps/api/internal/workerload/controller.go +++ b/apps/api/internal/workerload/controller.go @@ -226,7 +226,13 @@ func (l *Lease) EnterFinalizing(ctx context.Context) error { // Even under critical pressure, let one submitted task at a time finish // and release its provider lease instead of deadlocking the drain path. heavyLimit := max(c.heavyLimit, 1) - if c.preparing+c.finalizing < heavyLimit { + heavyInUseByOthers := c.preparing + c.finalizing + if l.phase == PhasePreparing { + // A preparing lease already owns one heavy slot. Moving that same + // lease to finalizing must not wait for itself to release the slot. + heavyInUseByOthers-- + } + if heavyInUseByOthers < heavyLimit { c.decrementPhaseLocked(l.phase) c.finalizing++ l.phase = PhaseFinalizing diff --git a/apps/api/internal/workerload/controller_test.go b/apps/api/internal/workerload/controller_test.go index 17d9de7..4d714d1 100644 --- a/apps/api/internal/workerload/controller_test.go +++ b/apps/api/internal/workerload/controller_test.go @@ -86,6 +86,23 @@ func TestWaitingReleasesHeavyPermitAndFinalizingReacquiresIt(t *testing.T) { } } +func TestPreparingLeaseCanTransitionDirectlyToFinalizingAtHeavyLimitOne(t *testing.T) { + controller := New(Config{Mode: ModeAdaptive, HardLimit: 1, InitialActive: 1, InitialHeavy: 1}) + lease, ok := controller.TryStart() + if !ok { + t.Fatal("preparing lease unavailable") + } + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + if err := lease.EnterFinalizing(ctx); err != nil { + t.Fatalf("preparing lease deadlocked while entering finalizing: %v", err) + } + if phase := lease.Phase(); phase != PhaseFinalizing { + t.Fatalf("lease phase=%s, want %s", phase, PhaseFinalizing) + } + lease.Release() +} + func TestLegacyModeUsesHardLimit(t *testing.T) { controller := New(Config{Mode: ModeLegacy, HardLimit: 7}) snapshot := controller.Observe(ResourceSample{MemoryCurrentBytes: 99, MemoryLimitBytes: 100, CPUUtilization: 1, CPUThrottled: true}) diff --git a/apps/api/migrations/0104_async_admission_dispatch_resilience.sql b/apps/api/migrations/0104_async_admission_dispatch_resilience.sql new file mode 100644 index 0000000..83382f7 --- /dev/null +++ b/apps/api/migrations/0104_async_admission_dispatch_resilience.sql @@ -0,0 +1,17 @@ +ALTER TABLE gateway_task_admissions + ADD COLUMN IF NOT EXISTS dispatch_next_at timestamptz NOT NULL DEFAULT now(), + ADD COLUMN IF NOT EXISTS dispatch_failure_count integer NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS dispatch_last_outcome text, + ADD COLUMN IF NOT EXISTS dispatch_last_error_code text, + ADD COLUMN IF NOT EXISTS dispatch_last_error_at timestamptz; + +ALTER TABLE gateway_task_admissions + ADD CONSTRAINT gateway_task_admissions_dispatch_failure_count_check + CHECK (dispatch_failure_count >= 0) NOT VALID; + +ALTER TABLE gateway_task_admissions + VALIDATE CONSTRAINT gateway_task_admissions_dispatch_failure_count_check; + +CREATE INDEX IF NOT EXISTS idx_task_admissions_async_dispatch_due + ON gateway_task_admissions(dispatch_next_at, priority, enqueued_at, task_id) + WHERE status = 'waiting' AND mode = 'async'; diff --git a/scripts/migration-safety-reviewed.json b/scripts/migration-safety-reviewed.json index 2b9bc89..8865a55 100644 --- a/scripts/migration-safety-reviewed.json +++ b/scripts/migration-safety-reviewed.json @@ -45,6 +45,13 @@ "DELETE FROM operation" ], "reason": "This one-time semantic migration takes a write-blocking/read-compatible lock, records count and SHA-256 manifests, archives every legacy allow row, verifies the archive, and only then deletes legacy allows. Deny rows are checksummed before cleanup and verified unchanged afterward. The procedural blocks make any count or checksum mismatch abort the transaction; reapplying after new whitelist writes also fails before deleting them." + }, + "apps/api/migrations/0104_async_admission_dispatch_resilience.sql": { + "sha256": "e4585aedc97a1f3e38ac6256863306cfbf8aca6dff623c97bc0f846eb3ff1767", + "allowedViolations": [ + "non-null column addition" + ], + "reason": "The dispatcher state columns use constant defaults on gateway_task_admissions, so PostgreSQL 18 can add them without rewriting the table and old binaries remain compatible. Production preflight on 2026-08-04 found 88 admission rows; the bounded nonnegative constraint validation and partial dispatch index are additive, and application rollback intentionally retains them." } } }