fix(queue): 隔离异步准入故障并持久化退避

将 X-Async 调度改为逐任务阻塞协议,区分全局容量、平台容量、用户组 FIFO、任务级异常和系统级故障,避免单个历史毒任务触发整批回滚。\n\n新增持久化退避、单任务 CAS 重选、幂等终态事务、固定标签指标和损坏快照自愈,并修复 Worker 从 preparing 直接进入 finalizing 时的自等待死锁。\n\n验证:API 全量测试、PostgreSQL 集成场景、race、go vet、govulncheck、pnpm lint/test/build、Compose 与发布脚本测试通过;pnpm audit 命中未改动的 Nx 工具链既有漏洞。
This commit is contained in:
2026-08-04 18:23:35 +08:00
parent 44d5cf2b9d
commit f4214dd489
12 changed files with 1654 additions and 133 deletions
@@ -96,6 +96,141 @@ SELECT
assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs)) 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) { func TestAsyncWorkerThousandConcurrentSchedulingAcceptance(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" { if databaseURL == "" {
@@ -537,7 +672,24 @@ func waitForTaskCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tas
} }
time.Sleep(100 * time.Millisecond) 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 { func activeLeaseExpiry(t *testing.T, ctx context.Context, pool *pgxpool.Pool, modelID string, wantCount int) time.Time {
+419 -53
View File
@@ -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 { if len(admissions) == 0 {
return false, nil return false, false, nil
} }
workerCapacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx) workerCapacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx)
if err != nil { if err != nil {
return false, err return false, false, err
} }
if workerCapacity <= 0 { if workerCapacity <= 0 {
return true, nil s.observeAsyncAdmissionDispatch("global_wait")
return true, false, nil
} }
progressed := false
microbatchSize := s.cfg.AsyncAdmissionMicrobatchSize microbatchSize := s.cfg.AsyncAdmissionMicrobatchSize
if microbatchSize <= 0 { if microbatchSize <= 0 {
microbatchSize = 1 microbatchSize = 1
@@ -765,7 +767,7 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []st
inputs := make([]store.TaskAdmissionInput, 0, end-start) inputs := make([]store.TaskAdmissionInput, 0, end-start)
tasksByID := make(map[string]store.GatewayTask, end-start) tasksByID := make(map[string]store.GatewayTask, end-start)
for _, admission := range admissions[start:end] { 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)) inputs = append(inputs, admissionInputFromSnapshot(admission, workerCapacity))
tasksByID[admission.TaskID] = store.GatewayTask{ID: admission.TaskID} tasksByID[admission.TaskID] = store.GatewayTask{ID: admission.TaskID}
continue continue
@@ -773,78 +775,411 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []st
task, loadErr := s.store.GetTask(ctx, admission.TaskID) task, loadErr := s.store.GetTask(ctx, admission.TaskID)
if loadErr != nil { if loadErr != nil {
if errors.Is(loadErr, pgx.ErrNoRows) { if errors.Is(loadErr, pgx.ErrNoRows) {
progressed = true
continue continue
} }
return false, loadErr return false, progressed, loadErr
} }
plan, planErr := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, authUserFromTask(task), &admission) plan, planErr := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, authUserFromTask(task), &admission)
if planErr != nil { 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 !plan.Eligible {
if enqueueErr := s.EnqueueAsyncTask(ctx, task); enqueueErr != nil { // This task already owns a durable admission row from its previous
return false, enqueueErr // 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, "") input := taskAdmissionInput(task, plan, "")
if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil && !errors.Is(rebindErr, pgx.ErrNoRows) { 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) inputs = append(inputs, input)
tasksByID[task.ID] = task tasksByID[task.ID] = task
} }
if len(inputs) == 0 { if len(inputs) == 0 {
continue 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, ctx,
inputs, inputs,
func(tx pgx.Tx, input store.TaskAdmissionInput) error { enqueueHook,
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))
},
) )
if err != nil { if batchErr != nil {
return false, err if isAsyncAdmissionSystemError(batchErr) {
} return false, progressed, batchErr
for _, outcome := range outcomes { }
s.observeTaskAdmissionAttempt(outcome.Result, outcome.Err) failedTaskID := asyncAdmissionFailedTaskID(outcomes)
if outcome.Err != nil { if failedTaskID == "" {
if errors.Is(outcome.Err, store.ErrTaskExecutionFinished) || return false, progressed, batchErr
errors.Is(outcome.Err, store.ErrQueueTimeout) { }
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 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 { continue
s.observeCandidateRouting("quota_race_rotated") }
platformModelID := outcome.Result.Admission.PlatformModelID for _, outcome := range outcomes {
if platformModelID == "" { input := inputByTaskID(inputs, outcome.TaskID)
for _, input := range inputs { globalSaturated, outcomeProgress, outcomeErr := s.handleAsyncAdmissionOutcome(
if input.TaskID == outcome.TaskID { ctx,
platformModelID = input.PlatformModelID outcome,
break input,
} tasksByID,
} admissionByTaskID(admissions[start:end], outcome.TaskID),
} )
if platformModelID != "" { if outcomeErr != nil {
marked, markErr := s.store.RequestWaitingTaskAdmissionReselect(ctx, platformModelID) return false, progressed, outcomeErr
if markErr != nil { }
return false, markErr progressed = progressed || outcomeProgress
} if globalSaturated {
if marked > 0 { return true, progressed, nil
s.observeTaskAdmission("candidate_reselect_requested")
}
}
return true, 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 { 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) cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
defer cancel() 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 { if s.logger != nil {
s.logger.Warn("cancel disconnected asynchronous task failed", "taskID", taskID, "error", err) 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) { func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
ticker := time.NewTicker(time.Second) ticker := time.NewTicker(time.Second)
defer ticker.Stop() defer ticker.Stop()
systemBackoff := time.Second
var systemRetryAt time.Time
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
@@ -873,21 +1219,34 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
case <-s.asyncAdmissionWake: case <-s.asyncAdmissionWake:
case <-ticker.C: case <-ticker.C:
} }
if time.Now().Before(systemRetryAt) {
continue
}
for { for {
batchLimit := asyncAdmissionBatchLimit(s.cfg.AsyncWorkerHardLimit) batchLimit := asyncAdmissionBatchLimit(s.cfg.AsyncWorkerHardLimit)
admissions, err := s.store.ListWaitingAsyncAdmissions(ctx, batchLimit) admissions, err := s.store.ListWaitingAsyncAdmissions(ctx, batchLimit)
if err != nil { if err != nil {
s.logger.Warn("list waiting async admissions failed", "error", err) 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 break
} }
if len(admissions) == 0 { if len(admissions) == 0 {
systemBackoff = time.Second
systemRetryAt = time.Time{}
break break
} }
saturated, err := s.dispatchWaitingAsyncTasks(ctx, admissions) saturated, progressed, err := s.dispatchWaitingAsyncTasks(ctx, admissions)
if err != nil { if err != nil {
s.logger.Warn("dispatch waiting async admission batch failed", "tasks", len(admissions), "error", err) 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 break
} }
systemBackoff = time.Second
systemRetryAt = time.Time{}
if saturated { if saturated {
// Every asynchronous task shares the global worker-capacity FIFO // Every asynchronous task shares the global worker-capacity FIFO
// scope. Once its current head cannot be admitted, later tasks // 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") s.observeTaskAdmission("dispatch_saturated")
break break
} }
if !progressed {
break
}
// PostgreSQL notifications are wake-up hints and collapse while a // PostgreSQL notifications are wake-up hints and collapse while a
// batch is being admitted. Keep filling consecutive batches until // batch is being admitted. Keep filling consecutive batches until
// the global execution scope is actually saturated so a large burst // the global execution scope is actually saturated so a large burst
@@ -919,7 +1281,11 @@ func (s *Service) reapExpiredTaskAdmissions(ctx context.Context) {
return return
case <-ticker.C: 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 err != nil {
if s.logger != nil { if s.logger != nil {
s.logger.Warn("reap expired task admissions failed", "error", err) s.logger.Warn("reap expired task admissions failed", "error", err)
@@ -1,12 +1,84 @@
package runner package runner
import ( import (
"errors"
"testing" "testing"
"time" "time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store" "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) { func TestAdmissionInputFromSnapshotRefreshesWorkerCapacity(t *testing.T) {
admission := store.TaskAdmission{ admission := store.TaskAdmission{
TaskID: "task-1", TaskID: "task-1",
@@ -80,6 +80,15 @@ type Metrics struct {
taskAdmissionCancelled atomic.Uint64 taskAdmissionCancelled atomic.Uint64
taskAdmissionExpired atomic.Uint64 taskAdmissionExpired atomic.Uint64
taskAdmissionMigrated 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 candidateNormalRotation atomic.Uint64
candidateFullAvoided atomic.Uint64 candidateFullAvoided atomic.Uint64
candidateQuotaRaceRotated 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) { func (m *Metrics) ObserveCandidateRouting(reason string) {
switch reason { switch reason {
case "normal_rotation": 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_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_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_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_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_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)) 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{ 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()}, {"expired", m.taskAdmissionExpired.Load()},
{"candidate_migrated", m.taskAdmissionMigrated.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) taskAdmissionWaitHistogram(w, m)
}) })
} }
@@ -62,6 +62,13 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
metrics.ObserveConcurrencyLeaseRenewal("lost") metrics.ObserveConcurrencyLeaseRenewal("lost")
metrics.ObserveTaskEventSkip("duplicate") metrics.ObserveTaskEventSkip("duplicate")
metrics.ObserveTaskEventSkip("budget_exceeded") 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() recorder := httptest.NewRecorder()
metrics.Handler(metricsSnapshot{ metrics.Handler(metricsSnapshot{
@@ -116,6 +123,13 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
`easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`, `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="duplicate"} 1`,
`easyai_gateway_task_events_skipped_total{outcome="budget_exceeded"} 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_max_connections 32`,
`easyai_gateway_postgres_pool_total_connections 18`, `easyai_gateway_postgres_pool_total_connections 18`,
`easyai_gateway_postgres_pool_acquired_connections 12`, `easyai_gateway_postgres_pool_acquired_connections 12`,
+321 -65
View File
@@ -65,21 +65,34 @@ type TaskAdmissionInput struct {
} }
type TaskAdmission struct { type TaskAdmission struct {
TaskID string TaskID string
PlatformID string PlatformID string
PlatformModelID string PlatformModelID string
UserGroupID string UserGroupID string
QueueKey string QueueKey string
Mode string Mode string
Status string Status string
Priority int Priority int
EnqueuedAt time.Time EnqueuedAt time.Time
WaitDeadlineAt time.Time WaitDeadlineAt time.Time
AdmittedAt time.Time AdmittedAt time.Time
WaiterID string WaiterID string
WaiterLeaseExpires time.Time WaiterLeaseExpires time.Time
Scopes []AdmissionScope Scopes []AdmissionScope
ReselectRequestedAt time.Time 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 { type TaskAdmissionResult struct {
@@ -87,6 +100,7 @@ type TaskAdmissionResult struct {
Admitted bool Admitted bool
NewlyAdmitted bool NewlyAdmitted bool
Leases []ConcurrencyLease Leases []ConcurrencyLease
Blockers []TaskAdmissionBlocker
} }
type TaskAdmissionBatchOutcome struct { type TaskAdmissionBatchOutcome struct {
@@ -96,12 +110,15 @@ type TaskAdmissionBatchOutcome struct {
} }
type TaskAdmissionMetricsSnapshot struct { type TaskAdmissionMetricsSnapshot struct {
QueueDepth int QueueDepth int
WaitingSync int WaitingSync int
WaitingAsync int WaitingAsync int
OldestWaitSeconds float64 DeferredAsync int
ExpiredWaiterBacklog int PendingReselect int
ExpiredDeadlineBacklog int OldestWaitSeconds float64
OldestDispatchableWaitSeconds float64
ExpiredWaiterBacklog int
ExpiredDeadlineBacklog int
} }
type TaskAdmissionReapResult struct { type TaskAdmissionReapResult struct {
@@ -347,7 +364,7 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
} }
return taskAdmissionTxOutcome{Result: result}, nil return taskAdmissionTxOutcome{Result: result}, nil
} }
if bindingChanged { if bindingChanged || activeLeases == 0 {
targetStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes) targetStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes)
if stateErr != nil { if stateErr != nil {
return taskAdmissionTxOutcome{}, stateErr 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 { if err != nil {
return taskAdmissionTxOutcome{}, err return taskAdmissionTxOutcome{}, err
} }
if !head { if !head {
return taskAdmissionTxOutcome{ return taskAdmissionTxOutcome{
Result: TaskAdmissionResult{Admission: admission}, Result: TaskAdmissionResult{Admission: admission, Blockers: headBlockers},
}, nil }, nil
} }
blockers := make([]TaskAdmissionBlocker, 0, len(scopeStates))
for _, state := range scopeStates { for _, state := range scopeStates {
if state.Saturated { if state.Saturated {
return taskAdmissionTxOutcome{ blockers = append(blockers, TaskAdmissionBlocker{
Result: TaskAdmissionResult{Admission: admission}, Reason: "saturated",
}, nil 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)) leases := make([]ConcurrencyLease, 0, len(scopes))
for _, scope := range scopes { for _, scope := range scopes {
@@ -498,9 +524,6 @@ func (s *Store) TryTaskAdmissionBatchWithAdmittedHook(
if admissionErr != nil { if admissionErr != nil {
return outcomes, admissionErr return outcomes, admissionErr
} }
if !result.Admitted {
break
}
} }
return outcomes, nil return outcomes, nil
} }
@@ -585,9 +608,6 @@ func (s *Store) tryTaskAdmissionAtomicBatchOnce(
if outcome.NotifyTaskID != "" { if outcome.NotifyTaskID != "" {
notify = true notify = true
} }
if !outcome.Result.Admitted {
break
}
} }
if err := tx.Commit(ctx); err != nil { if err := tx.Commit(ctx); err != nil {
return outcomes, err return outcomes, err
@@ -713,6 +733,11 @@ SET platform_id = $2::uuid,
wait_deadline_at = LEAST(wait_deadline_at, $6), wait_deadline_at = LEAST(wait_deadline_at, $6),
scope_snapshot = $7::jsonb, scope_snapshot = $7::jsonb,
reselect_requested_at = NULL, 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() updated_at = now()
WHERE task_id = $1::uuid WHERE task_id = $1::uuid
AND status = 'waiting' 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, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority,
enqueued_at, wait_deadline_at, admitted_at, enqueued_at, wait_deadline_at, admitted_at,
COALESCE(waiter_id, ''), waiter_lease_expires_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.TaskID,
input.PlatformID, input.PlatformID,
input.PlatformModelID, input.PlatformModelID,
@@ -779,13 +806,20 @@ SET status = 'waiting',
waiter_id = NULLIF($7, ''), waiter_id = NULLIF($7, ''),
waiter_lease_expires_at = $8, waiter_lease_expires_at = $8,
scope_snapshot = $9::jsonb, 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() updated_at = now()
WHERE task_id = $1::uuid WHERE task_id = $1::uuid
RETURNING task_id::text, platform_id::text, platform_model_id::text, RETURNING task_id::text, platform_id::text, platform_model_id::text,
COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority,
enqueued_at, wait_deadline_at, admitted_at, enqueued_at, wait_deadline_at, admitted_at,
COALESCE(waiter_id, ''), waiter_lease_expires_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.TaskID,
input.PlatformID, input.PlatformID,
input.PlatformModelID, input.PlatformModelID,
@@ -915,6 +949,7 @@ LEFT JOIN river_job job ON job.id = task.river_job_id
WHERE admission.mode = 'async' WHERE admission.mode = 'async'
AND task.status = 'queued' AND task.status = 'queued'
AND task.next_run_at <= now() AND task.next_run_at <= now()
AND admission.dispatch_next_at <= now()
AND ( AND (
admission.status = 'waiting' admission.status = 'waiting'
OR ( 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, COALESCE(admission.user_group_id::text, ''), admission.queue_key, admission.mode,
admission.status, admission.priority, admission.enqueued_at, admission.wait_deadline_at, admission.status, admission.priority, admission.enqueued_at, admission.wait_deadline_at,
admission.admitted_at, COALESCE(admission.waiter_id, ''), admission.waiter_lease_expires_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 FROM gateway_task_admissions admission
JOIN gateway_tasks task ON task.id = admission.task_id JOIN gateway_tasks task ON task.id = admission.task_id
LEFT JOIN river_job job ON job.id = task.river_job_id LEFT JOIN river_job job ON job.id = task.river_job_id
WHERE admission.mode = 'async' WHERE admission.mode = 'async'
AND task.status = 'queued' AND task.status = 'queued'
AND task.next_run_at <= now() AND task.next_run_at <= now()
AND admission.dispatch_next_at <= now()
AND ( AND (
admission.status = 'waiting' admission.status = 'waiting'
OR ( OR (
@@ -1006,27 +1044,105 @@ LIMIT $1`, limit)
return admissions, rows.Err() return admissions, rows.Err()
} }
// RequestWaitingTaskAdmissionReselect marks every queued task bound to a // RequestTaskAdmissionReselect marks one waiting task for candidate
// saturated platform model so the dispatcher can route the next batch to a // reselection. The expected platform-model binding makes the operation a CAS:
// different eligible candidate. The durable marker lets multiple dispatchers // a concurrent dispatcher that already migrated or admitted the task wins.
// observe the same decision without assigning tasks to a specific Worker. func (s *Store) RequestTaskAdmissionReselect(
func (s *Store) RequestWaitingTaskAdmissionReselect(ctx context.Context, platformModelID string) (int64, error) { ctx context.Context,
result, err := s.pool.Exec(ctx, ` taskID string,
expectedPlatformModelID string,
) (TaskAdmission, bool, error) {
row := s.pool.QueryRow(ctx, `
UPDATE gateway_task_admissions admission UPDATE gateway_task_admissions admission
SET reselect_requested_at = now(), 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() updated_at = now()
FROM gateway_tasks task FROM gateway_tasks task
WHERE admission.task_id = task.id WHERE admission.task_id = $1::uuid
AND admission.platform_model_id = $1::uuid AND admission.task_id = task.id
AND admission.platform_model_id = $2::uuid
AND admission.mode = 'async' AND admission.mode = 'async'
AND admission.status = 'waiting' AND admission.status = 'waiting'
AND task.status = 'queued' AND task.status = 'queued'
AND task.next_run_at <= now() AND task.next_run_at <= now()
AND admission.reselect_requested_at IS NULL`, platformModelID) AND admission.reselect_requested_at IS NULL
if err != nil { RETURNING admission.task_id::text, admission.platform_id::text, admission.platform_model_id::text,
return 0, err 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 // 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, SELECT COUNT(*) FILTER (WHERE status = 'waiting')::int,
COUNT(*) FILTER (WHERE status = 'waiting' AND mode = 'sync')::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')::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')), 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 ( COUNT(*) FILTER (
WHERE status = 'waiting' AND mode = 'sync' AND waiter_lease_expires_at <= now() WHERE status = 'waiting' AND mode = 'sync' AND waiter_lease_expires_at <= now()
)::int, )::int,
@@ -1105,17 +1230,28 @@ FROM gateway_task_admissions`).Scan(
&snapshot.QueueDepth, &snapshot.QueueDepth,
&snapshot.WaitingSync, &snapshot.WaitingSync,
&snapshot.WaitingAsync, &snapshot.WaitingAsync,
&snapshot.DeferredAsync,
&snapshot.PendingReselect,
&snapshot.OldestWaitSeconds, &snapshot.OldestWaitSeconds,
&snapshot.OldestDispatchableWaitSeconds,
&snapshot.ExpiredWaiterBacklog, &snapshot.ExpiredWaiterBacklog,
&snapshot.ExpiredDeadlineBacklog, &snapshot.ExpiredDeadlineBacklog,
) )
return snapshot, err 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 { if limit <= 0 || limit > 1000 {
limit = 500 limit = 500
} }
defaultCallbackURL := ""
if len(defaultCallbackURLs) > 0 {
defaultCallbackURL = strings.TrimSpace(defaultCallbackURLs[0])
}
tx, err := s.pool.Begin(ctx) tx, err := s.pool.Begin(ctx)
if err != nil { if err != nil {
return TaskAdmissionReapResult{}, err return TaskAdmissionReapResult{}, err
@@ -1155,6 +1291,7 @@ LIMIT $1`, limit)
return TaskAdmissionReapResult{}, err return TaskAdmissionReapResult{}, err
} }
var deadlineExpired bool var deadlineExpired bool
var dispatchFailureCount int
err := tx.QueryRow(ctx, ` err := tx.QueryRow(ctx, `
DELETE FROM gateway_task_admissions DELETE FROM gateway_task_admissions
WHERE task_id = $1::uuid WHERE task_id = $1::uuid
@@ -1163,7 +1300,7 @@ WHERE task_id = $1::uuid
wait_deadline_at <= now() wait_deadline_at <= now()
OR (mode = 'sync' AND waiter_lease_expires_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) { if errors.Is(err, pgx.ErrNoRows) {
continue continue
} }
@@ -1177,22 +1314,112 @@ RETURNING wait_deadline_at <= now()`, taskID).Scan(&deadlineExpired)
code = "queue_timeout" code = "queue_timeout"
message = ErrQueueTimeout.Error() message = ErrQueueTimeout.Error()
status = "failed" status = "failed"
if dispatchFailureCount > 0 {
code = "admission_dispatch_failed"
message = "task admission dispatcher could not prepare the task before its queue deadline"
}
result.ExpiredDeadlines++ result.ExpiredDeadlines++
} else { } else {
result.ExpiredWaiters++ result.ExpiredWaiters++
} }
if _, err := tx.Exec(ctx, ` publicErrorJSON := encodePublicErrorSnapshot(code, message, 0, true, "", taskID)
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks UPDATE gateway_tasks
SET status = $2, SET status = $2,
error = $3, error = NULL,
error_code = $4, error_code = $4,
error_message = $3, 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(), finished_at = now(),
updated_at = now() updated_at = now()
WHERE id = $1::uuid 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 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++ reaped++
} }
if err := tx.Commit(ctx); err != nil { 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, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority,
enqueued_at, wait_deadline_at, admitted_at, enqueued_at, wait_deadline_at, admitted_at,
COALESCE(waiter_id, ''), waiter_lease_expires_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.TaskID, input.PlatformID, input.PlatformModelID, input.UserGroupID,
input.QueueKey, input.Mode, input.Priority, deadline, waiterID, waiterLease, scopeSnapshot) input.QueueKey, input.Mode, input.Priority, deadline, waiterID, waiterLease, scopeSnapshot)
return scanTaskAdmission(row) 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, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority,
enqueued_at, wait_deadline_at, admitted_at, enqueued_at, wait_deadline_at, admitted_at,
COALESCE(waiter_id, ''), waiter_lease_expires_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()) taskID, waiterID, admissionWaiterLeaseTTL.String())
return scanTaskAdmission(row) return scanTaskAdmission(row)
} }
@@ -1478,13 +1709,19 @@ UPDATE gateway_task_admissions
SET status = 'admitted', SET status = 'admitted',
admitted_at = now(), admitted_at = now(),
waiter_lease_expires_at = NULL, 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() updated_at = now()
WHERE task_id = $1::uuid AND status = 'waiting' WHERE task_id = $1::uuid AND status = 'waiting'
RETURNING task_id::text, platform_id::text, platform_model_id::text, RETURNING task_id::text, platform_id::text, platform_model_id::text,
COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority,
enqueued_at, wait_deadline_at, admitted_at, enqueued_at, wait_deadline_at, admitted_at,
COALESCE(waiter_id, ''), waiter_lease_expires_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) taskID)
return scanTaskAdmission(row) 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, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority,
enqueued_at, wait_deadline_at, admitted_at, enqueued_at, wait_deadline_at, admitted_at,
COALESCE(waiter_id, ''), waiter_lease_expires_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 FROM gateway_task_admissions
WHERE task_id = $1::uuid`, taskID) WHERE task_id = $1::uuid`, taskID)
admission, err := scanTaskAdmission(row) admission, err := scanTaskAdmission(row)
@@ -1518,6 +1757,7 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) {
var admittedAt sql.NullTime var admittedAt sql.NullTime
var waiterLeaseExpires sql.NullTime var waiterLeaseExpires sql.NullTime
var reselectRequestedAt sql.NullTime var reselectRequestedAt sql.NullTime
var dispatchLastErrorAt sql.NullTime
var scopeSnapshot []byte var scopeSnapshot []byte
err := row.Scan( err := row.Scan(
&admission.TaskID, &admission.TaskID,
@@ -1535,6 +1775,11 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) {
&waiterLeaseExpires, &waiterLeaseExpires,
&scopeSnapshot, &scopeSnapshot,
&reselectRequestedAt, &reselectRequestedAt,
&admission.DispatchNextAt,
&admission.DispatchFailureCount,
&admission.DispatchLastOutcome,
&admission.DispatchLastErrorCode,
&dispatchLastErrorAt,
) )
if admittedAt.Valid { if admittedAt.Valid {
admission.AdmittedAt = admittedAt.Time admission.AdmittedAt = admittedAt.Time
@@ -1545,15 +1790,18 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) {
if reselectRequestedAt.Valid { if reselectRequestedAt.Valid {
admission.ReselectRequestedAt = reselectRequestedAt.Time admission.ReselectRequestedAt = reselectRequestedAt.Time
} }
if dispatchLastErrorAt.Valid {
admission.DispatchLastErrorAt = dispatchLastErrorAt.Time
}
if err == nil && len(scopeSnapshot) > 0 { if err == nil && len(scopeSnapshot) > 0 {
if unmarshalErr := json.Unmarshal(scopeSnapshot, &admission.Scopes); unmarshalErr != nil { 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 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 { for _, scope := range scopes {
var head string var head string
switch scope.ScopeType { switch scope.ScopeType {
@@ -1561,10 +1809,12 @@ func isTaskAdmissionHeadTx(ctx context.Context, tx pgx.Tx, admission TaskAdmissi
if err := tx.QueryRow(ctx, ` if err := tx.QueryRow(ctx, `
SELECT task_id::text SELECT task_id::text
FROM gateway_task_admissions 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 ORDER BY priority ASC, enqueued_at ASC, task_id ASC
LIMIT 1`, admission.PlatformModelID).Scan(&head); err != nil { LIMIT 1`, admission.PlatformModelID).Scan(&head); err != nil {
return false, err return false, nil, err
} }
case "user_group": case "user_group":
if admission.UserGroupID == "" { if admission.UserGroupID == "" {
@@ -1573,19 +1823,25 @@ LIMIT 1`, admission.PlatformModelID).Scan(&head); err != nil {
if err := tx.QueryRow(ctx, ` if err := tx.QueryRow(ctx, `
SELECT task_id::text SELECT task_id::text
FROM gateway_task_admissions 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 ORDER BY priority ASC, enqueued_at ASC, task_id ASC
LIMIT 1`, admission.UserGroupID).Scan(&head); err != nil { LIMIT 1`, admission.UserGroupID).Scan(&head); err != nil {
return false, err return false, nil, err
} }
default: default:
continue continue
} }
if head != admission.TaskID { 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 { func expireTaskAdmissionTx(ctx context.Context, tx pgx.Tx, taskID string) error {
@@ -131,6 +131,9 @@ RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID);
if err != nil || result.Admitted { if err != nil || result.Admitted {
t.Fatalf("second task should wait: result=%+v err=%v", result, err) 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) higherPriorityAsync := createTask(true)
result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, "")) result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
if err != nil || result.Admitted { if err != nil || result.Admitted {
@@ -176,6 +179,9 @@ WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil {
if result.Admitted { if result.Admitted {
t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter") 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, "")) result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
if err != nil || !result.Admitted { if err != nil || !result.Admitted {
t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err) 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) { if listedSnapshot == nil || len(listedSnapshot.Scopes) != len(queuedAdmission.Scopes) {
t.Fatalf("listed admission snapshot=%+v, want %d scopes", listedSnapshot, 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 { if err != nil {
t.Fatalf("request waiting admission reselection: %v", err) t.Fatalf("request waiting admission reselection: %v", err)
} }
if markedForReselect < 1 { if !markedForReselect {
t.Fatalf("marked admissions=%d, want at least the queued task", 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) reselectAdmission, err := first.GetTaskAdmission(ctx, queuedAtomicTask.ID)
if err != nil { 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) 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) terminalResidue := createTask(true)
result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, "")) result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, ""))
if err != nil || !result.Admitted { 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] { if !leaderSet[tasks[0].ID] || leaderSet[tasks[2].ID] {
t.Fatalf("grouped candidate leaders = %v, want only group head %s", leaders, tasks[0].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) { func TestWorkerCapacityAllocationAndFailover(t *testing.T) {
+147 -10
View File
@@ -545,17 +545,40 @@ WHERE id = $1::uuid
} }
func (s *Store) FailQueuedTask(ctx context.Context, taskID string, code string, message string) (GatewayTask, error) { 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) 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 UPDATE gateway_tasks
SET status = 'failed', SET status = 'failed',
error = NULL, error = NULL,
error_code = NULLIF($2, ''), error_code = NULLIF($2, ''),
error_message = NULLIF($3, ''), error_message = NULLIF($3, ''),
public_error = $4::jsonb, public_error = $4::jsonb,
billing_status = CASE billing_status = CASE
WHEN run_mode NOT IN ('production', 'acceptance', 'acceptance_canary') OR gateway_user_id IS NULL THEN 'not_required' 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, END,
billing_updated_at = now(), billing_updated_at = now(),
locked_by = NULL, locked_by = NULL,
@@ -566,14 +589,76 @@ SET status = 'failed',
finished_at = now(), finished_at = now(),
updated_at = now() updated_at = now()
WHERE id = $1::uuid 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 { if err != nil {
return GatewayTask{}, err return GatewayTask{}, err
} }
if tag.RowsAffected() != 1 { s.notifyTaskAdmissionBestEffort(ctx, "*")
return GatewayTask{}, ErrTaskExecutionFinished return task, nil
}
return s.GetTask(ctx, taskID)
} }
func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error { 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) { 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) message = strings.TrimSpace(message)
if message == "" { if message == "" {
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 { if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
return err return err
} }
if _, err := releaseRuntimeRateReservationsTx(ctx, tx, []string{taskID}); err != nil {
return err
}
payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "queued_cancelled"}) payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "queued_cancelled"})
_, err = tx.Exec(ctx, ` if _, err = tx.Exec(ctx, `
INSERT INTO settlement_outbox ( INSERT INTO settlement_outbox (
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at 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 run_mode IN ('production', 'acceptance', 'acceptance_canary')
AND gateway_user_id IS NOT NULL AND gateway_user_id IS NOT NULL
AND reservation_amount > 0 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 return err
}) })
if err != nil { if err != nil {
+7 -1
View File
@@ -226,7 +226,13 @@ func (l *Lease) EnterFinalizing(ctx context.Context) error {
// Even under critical pressure, let one submitted task at a time finish // Even under critical pressure, let one submitted task at a time finish
// and release its provider lease instead of deadlocking the drain path. // and release its provider lease instead of deadlocking the drain path.
heavyLimit := max(c.heavyLimit, 1) 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.decrementPhaseLocked(l.phase)
c.finalizing++ c.finalizing++
l.phase = PhaseFinalizing l.phase = PhaseFinalizing
@@ -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) { func TestLegacyModeUsesHardLimit(t *testing.T) {
controller := New(Config{Mode: ModeLegacy, HardLimit: 7}) controller := New(Config{Mode: ModeLegacy, HardLimit: 7})
snapshot := controller.Observe(ResourceSample{MemoryCurrentBytes: 99, MemoryLimitBytes: 100, CPUUtilization: 1, CPUThrottled: true}) snapshot := controller.Observe(ResourceSample{MemoryCurrentBytes: 99, MemoryLimitBytes: 100, CPUUtilization: 1, CPUThrottled: true})
@@ -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';
+7
View File
@@ -45,6 +45,13 @@
"DELETE FROM operation" "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." "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."
} }
} }
} }