package store import ( "context" "database/sql" "errors" "fmt" "math" "sort" "strings" "time" "github.com/jackc/pgx/v5" ) const ( admissionWaiterLeaseTTL = 15 * time.Second admissionNotifyChannel = "gateway_task_admission" executionSlotWaitMax = 24 * time.Hour ) var ErrQueueTimeout = errors.New("task admission queue wait timed out") type QueueTimeoutError struct { TaskID string } func (e *QueueTimeoutError) Error() string { return ErrQueueTimeout.Error() } func (e *QueueTimeoutError) ErrorCode() string { return "queue_timeout" } func (e *QueueTimeoutError) Unwrap() error { return ErrQueueTimeout } type AdmissionScope struct { ScopeType string ScopeKey string ScopeName string ScopeMetadata map[string]any ConcurrentLimit float64 Amount float64 LeaseTTLSeconds int QueueLimit int MaxWaitSeconds int Policy map[string]any } type TaskAdmissionInput struct { TaskID string PlatformID string PlatformModelID string UserGroupID string QueueKey string Mode string Priority int WaiterID string Scopes []AdmissionScope } type TaskAdmission struct { TaskID string PlatformID string PlatformModelID string UserGroupID string QueueKey string Mode string Status string Priority int EnqueuedAt time.Time WaitDeadlineAt time.Time AdmittedAt time.Time WaiterID string WaiterLeaseExpires time.Time } type TaskAdmissionResult struct { Admission TaskAdmission Admitted bool NewlyAdmitted bool Leases []ConcurrencyLease } type TaskAdmissionMetricsSnapshot struct { QueueDepth int WaitingSync int WaitingAsync int OldestWaitSeconds float64 ExpiredWaiterBacklog int ExpiredDeadlineBacklog int } type TaskAdmissionReapResult struct { ExpiredWaiters int ExpiredDeadlines int } func (s *Store) TryTaskAdmission(ctx context.Context, input TaskAdmissionInput) (TaskAdmissionResult, error) { return s.tryTaskAdmission(ctx, input, nil) } // TryTaskAdmissionWithAdmittedHook runs onAdmitted in the same PostgreSQL // transaction that changes the admission to admitted. Asynchronous callers use // this hook to insert the unique River job without leaving a crash window // between capacity admission and durable job creation. func (s *Store) TryTaskAdmissionWithAdmittedHook( ctx context.Context, input TaskAdmissionInput, onAdmitted func(pgx.Tx) error, ) (TaskAdmissionResult, error) { return s.tryTaskAdmission(ctx, input, onAdmitted) } // QueueTaskAdmissionWithHook durably registers an asynchronous FIFO waiter and // runs onQueued in the same transaction. It intentionally does not reserve // concurrency leases: the Worker dispatcher obtains business and execution // capacity only when the task reaches the head of every applicable queue. func (s *Store) QueueTaskAdmissionWithHook( ctx context.Context, input TaskAdmissionInput, onQueued func(pgx.Tx) error, ) (TaskAdmission, error) { if err := validateTaskAdmissionInput(input); err != nil { return TaskAdmission{}, err } if input.Mode != "async" { return TaskAdmission{}, errors.New("queued admission registration requires asynchronous mode") } tx, err := s.pool.Begin(ctx) if err != nil { return TaskAdmission{}, err } defer tx.Rollback(ctx) if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil { return TaskAdmission{}, err } var taskActive bool if err := tx.QueryRow(ctx, ` SELECT status IN ('queued', 'running') FROM gateway_tasks WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { return TaskAdmission{}, err } if !taskActive { return TaskAdmission{}, ErrTaskExecutionFinished } admission, found, err := loadTaskAdmissionTx(ctx, tx, input.TaskID) if err != nil { return TaskAdmission{}, err } scopes := normalizedAdmissionScopes(input.Scopes) lockScopes := append([]AdmissionScope{}, scopes...) if found { lockScopes = append(lockScopes, AdmissionScope{ScopeType: "platform_model", ScopeKey: admission.PlatformModelID}, AdmissionScope{ScopeType: "user_group", ScopeKey: admission.UserGroupID}, ) } for _, scope := range normalizedAdmissionScopes(lockScopes) { if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil { return TaskAdmission{}, err } } if found && admission.Status == "admitted" { activeLeases, leaseErr := activeTaskAdmissionLeaseCountTx(ctx, tx, input.TaskID) if leaseErr != nil { return TaskAdmission{}, leaseErr } if activeLeases == 0 { admission, err = resetTaskAdmissionToWaitingTx(ctx, tx, input) if err != nil { return TaskAdmission{}, err } } } if found && admission.Status == "waiting" && !admission.WaitDeadlineAt.After(time.Now()) { if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { return TaskAdmission{}, err } if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return TaskAdmission{}, err } if err := tx.Commit(ctx); err != nil { return TaskAdmission{}, err } return TaskAdmission{}, &QueueTimeoutError{TaskID: input.TaskID} } if !found { scopeStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes) if stateErr != nil { return TaskAdmission{}, stateErr } if admissionErr := validateQueuedAdmissionCapacity(scopeStates); admissionErr != nil { return TaskAdmission{}, admissionErr } admission, err = insertTaskAdmissionTx(ctx, tx, input, admissionDeadline(scopes)) if err != nil { return TaskAdmission{}, err } found = true } if onQueued != nil { if err := onQueued(tx); err != nil { return TaskAdmission{}, err } } if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { return TaskAdmission{}, err } if err := tx.Commit(ctx); err != nil { return TaskAdmission{}, err } return admission, nil } func (s *Store) tryTaskAdmission( ctx context.Context, input TaskAdmissionInput, onAdmitted func(pgx.Tx) error, ) (TaskAdmissionResult, error) { if err := validateTaskAdmissionInput(input); err != nil { return TaskAdmissionResult{}, err } tx, err := s.pool.Begin(ctx) if err != nil { return TaskAdmissionResult{}, err } defer tx.Rollback(ctx) if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil { return TaskAdmissionResult{}, err } var taskActive bool if err := tx.QueryRow(ctx, ` SELECT status IN ('queued', 'running') FROM gateway_tasks WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { return TaskAdmissionResult{}, err } if !taskActive { return TaskAdmissionResult{}, ErrTaskExecutionFinished } admission, found, err := loadTaskAdmissionTx(ctx, tx, input.TaskID) if err != nil { return TaskAdmissionResult{}, err } scopes := normalizedAdmissionScopes(input.Scopes) lockScopes := append([]AdmissionScope{}, scopes...) if found { lockScopes = append(lockScopes, AdmissionScope{ScopeType: "platform_model", ScopeKey: admission.PlatformModelID}, AdmissionScope{ScopeType: "user_group", ScopeKey: admission.UserGroupID}, ) } for _, scope := range normalizedAdmissionScopes(lockScopes) { if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil { return TaskAdmissionResult{}, err } } if found && admission.Status == "admitted" { activeLeases, leaseErr := activeTaskAdmissionLeaseCountTx(ctx, tx, input.TaskID) if leaseErr != nil { return TaskAdmissionResult{}, leaseErr } bindingChanged := admission.PlatformID != input.PlatformID || admission.PlatformModelID != input.PlatformModelID || admission.UserGroupID != input.UserGroupID if activeLeases > 0 && !bindingChanged { result := TaskAdmissionResult{Admission: admission, Admitted: true} if onAdmitted != nil { if err := onAdmitted(tx); err != nil { return TaskAdmissionResult{}, err } } if err := tx.Commit(ctx); err != nil { return TaskAdmissionResult{}, err } return result, nil } if bindingChanged { targetStates, stateErr := admissionScopeStatesTx(ctx, tx, scopes) if stateErr != nil { return TaskAdmissionResult{}, stateErr } if admissionErr := validateNewAdmissionCapacity(targetStates); admissionErr != nil { return TaskAdmissionResult{}, admissionErr } } admission, err = resetTaskAdmissionToWaitingTx(ctx, tx, input) if err != nil { return TaskAdmissionResult{}, err } } if found && !admission.WaitDeadlineAt.After(time.Now()) { if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { return TaskAdmissionResult{}, err } if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return TaskAdmissionResult{}, err } if err := tx.Commit(ctx); err != nil { return TaskAdmissionResult{}, err } return TaskAdmissionResult{}, &QueueTimeoutError{TaskID: input.TaskID} } scopeStates, err := admissionScopeStatesTx(ctx, tx, scopes) if err != nil { return TaskAdmissionResult{}, err } if !found { if admissionErr := validateNewAdmissionCapacity(scopeStates); admissionErr != nil { return TaskAdmissionResult{}, admissionErr } deadline := admissionDeadline(scopes) admission, err = insertTaskAdmissionTx(ctx, tx, input, deadline) if err != nil { return TaskAdmissionResult{}, err } found = true } else if input.Mode == "sync" && strings.TrimSpace(input.WaiterID) != "" { admission, err = renewTaskAdmissionWaiterTx(ctx, tx, input.TaskID, input.WaiterID) if err != nil { return TaskAdmissionResult{}, err } } head, err := isTaskAdmissionHeadTx(ctx, tx, admission, scopes) if err != nil { return TaskAdmissionResult{}, err } if !head { if err := tx.Commit(ctx); err != nil { return TaskAdmissionResult{}, err } return TaskAdmissionResult{Admission: admission}, nil } for _, state := range scopeStates { if state.Saturated { if err := tx.Commit(ctx); err != nil { return TaskAdmissionResult{}, err } return TaskAdmissionResult{Admission: admission}, nil } } leases := make([]ConcurrencyLease, 0, len(scopes)) for _, scope := range scopes { if scope.ConcurrentLimit <= 0 || scope.Amount <= 0 { continue } lease, err := reserveConcurrencyLease(ctx, tx, input.TaskID, "", RateLimitReservation{ ScopeType: scope.ScopeType, ScopeKey: scope.ScopeKey, ScopeName: scope.ScopeName, ScopeMetadata: scope.ScopeMetadata, Metric: "concurrent", Limit: scope.ConcurrentLimit, Amount: scope.Amount, LeaseTTLSeconds: scope.LeaseTTLSeconds, Policy: scope.Policy, }) if err != nil { if errors.Is(err, ErrRateLimited) { if err := tx.Commit(ctx); err != nil { return TaskAdmissionResult{}, err } return TaskAdmissionResult{Admission: admission}, nil } return TaskAdmissionResult{}, err } leases = append(leases, lease) } admission, err = markTaskAdmissionAdmittedTx(ctx, tx, input.TaskID) if err != nil { return TaskAdmissionResult{}, err } // Continue the FIFO chain after one task is admitted. Every API process // only probes the global head it owns, so free capacity is filled without // broadcasting an advisory-lock attempt to every waiter. if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return TaskAdmissionResult{}, err } if onAdmitted != nil { if err := onAdmitted(tx); err != nil { return TaskAdmissionResult{}, err } } if err := tx.Commit(ctx); err != nil { return TaskAdmissionResult{}, err } return TaskAdmissionResult{Admission: admission, Admitted: true, NewlyAdmitted: true, Leases: leases}, nil } func validateNewAdmissionCapacity(states []admissionScopeState) error { mustWait := false for _, state := range states { if !state.Saturated { continue } mustWait = true if state.Scope.QueueLimit <= 0 { return concurrencyAdmissionError(state) } } if !mustWait { return nil } for _, state := range states { if state.Scope.QueueLimit > 0 && state.Waiting >= state.Scope.QueueLimit { return queueFullAdmissionError(state) } } return nil } func validateQueuedAdmissionCapacity(states []admissionScopeState) error { // Waiting asynchronous tasks have not acquired leases yet. Count them as // virtual capacity reservations so a burst cannot bypass concurrent and // queue limits before the dispatcher creates River jobs. for _, state := range states { if state.Scope.ConcurrentLimit <= 0 { continue } virtualUsage := state.Active + float64(state.Waiting)*state.Scope.Amount + state.Scope.Amount if virtualUsage <= state.Scope.ConcurrentLimit { continue } if state.Scope.QueueLimit <= 0 { return concurrencyAdmissionError(state) } virtualQueueDepth := int(math.Ceil(virtualUsage - state.Scope.ConcurrentLimit)) if virtualQueueDepth > state.Scope.QueueLimit { return queueFullAdmissionError(state) } } return nil } // RebindWaitingTaskAdmission migrates a queued task to a newly selected // candidate without changing its FIFO age or priority. Capacity checks happen // while both the old and new queue locks are held. func (s *Store) RebindWaitingTaskAdmission(ctx context.Context, input TaskAdmissionInput) (TaskAdmission, error) { if err := validateTaskAdmissionInput(input); err != nil { return TaskAdmission{}, err } tx, err := s.pool.Begin(ctx) if err != nil { return TaskAdmission{}, err } defer tx.Rollback(ctx) if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+input.TaskID); err != nil { return TaskAdmission{}, err } current, found, err := loadTaskAdmissionTx(ctx, tx, input.TaskID) if err != nil { return TaskAdmission{}, err } if !found { return TaskAdmission{}, pgx.ErrNoRows } if current.Status != "waiting" { return current, nil } scopes := normalizedAdmissionScopes(input.Scopes) lockScopes := append([]AdmissionScope{}, scopes...) lockScopes = append(lockScopes, AdmissionScope{ScopeType: "platform_model", ScopeKey: current.PlatformModelID}, AdmissionScope{ScopeType: "user_group", ScopeKey: current.UserGroupID}, ) for _, scope := range normalizedAdmissionScopes(lockScopes) { if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, admissionLockKey(scope)); err != nil { return TaskAdmission{}, err } } states, err := admissionScopeStatesTx(ctx, tx, scopes) if err != nil { return TaskAdmission{}, err } for index := range states { sameScope := (states[index].Scope.ScopeType == "platform_model" && states[index].Scope.ScopeKey == current.PlatformModelID) || (states[index].Scope.ScopeType == "user_group" && states[index].Scope.ScopeKey == current.UserGroupID) if sameScope && states[index].Waiting > 0 { states[index].Waiting-- } } if admissionErr := validateNewAdmissionCapacity(states); admissionErr != nil { return TaskAdmission{}, admissionErr } deadline := admissionDeadline(scopes) row := tx.QueryRow(ctx, ` UPDATE gateway_task_admissions SET platform_id = $2::uuid, platform_model_id = $3::uuid, user_group_id = NULLIF($4, '')::uuid, queue_key = $5, wait_deadline_at = LEAST(wait_deadline_at, $6), reselect_requested_at = NULL, updated_at = now() WHERE task_id = $1::uuid AND status = 'waiting' RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at`, input.TaskID, input.PlatformID, input.PlatformModelID, input.UserGroupID, input.QueueKey, deadline, ) migrated, err := scanTaskAdmission(row) if err != nil { return TaskAdmission{}, err } if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil { return TaskAdmission{}, err } if err := tx.Commit(ctx); err != nil { return TaskAdmission{}, err } return migrated, nil } func activeTaskAdmissionLeaseCountTx(ctx context.Context, tx pgx.Tx, taskID string) (int, error) { var count int err := tx.QueryRow(ctx, ` SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL AND expires_at > now()`, taskID).Scan(&count) return count, err } func resetTaskAdmissionToWaitingTx(ctx context.Context, tx pgx.Tx, input TaskAdmissionInput) (TaskAdmission, error) { if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases SET released_at = now() WHERE task_id = $1::uuid AND released_at IS NULL`, input.TaskID); err != nil { return TaskAdmission{}, err } var waiterLease any if input.Mode == "sync" { waiterLease = time.Now().Add(admissionWaiterLeaseTTL) } row := tx.QueryRow(ctx, ` UPDATE gateway_task_admissions SET status = 'waiting', admitted_at = NULL, platform_id = $2::uuid, platform_model_id = $3::uuid, user_group_id = NULLIF($4, '')::uuid, queue_key = $5, wait_deadline_at = LEAST(wait_deadline_at, $6), waiter_id = NULLIF($7, ''), waiter_lease_expires_at = $8, updated_at = now() WHERE task_id = $1::uuid RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at`, input.TaskID, input.PlatformID, input.PlatformModelID, input.UserGroupID, input.QueueKey, admissionDeadline(input.Scopes), strings.TrimSpace(input.WaiterID), waiterLease, ) return scanTaskAdmission(row) } func (s *Store) RenewTaskAdmissionWaiter(ctx context.Context, taskID string, waiterID string) error { result, err := s.pool.Exec(ctx, ` UPDATE gateway_task_admissions SET waiter_lease_expires_at = now() + $3::interval, updated_at = now() WHERE task_id = $1::uuid AND mode = 'sync' AND status = 'waiting' AND waiter_id = $2`, taskID, waiterID, admissionWaiterLeaseTTL.String()) if err != nil { return err } if result.RowsAffected() == 0 { return pgx.ErrNoRows } return nil } // RenewTaskAdmissionWaiters refreshes all synchronous waiters owned by one API // process in a single statement. This keeps the 15-second crash-recovery lease // without running a full advisory-lock admission transaction per waiter. func (s *Store) RenewTaskAdmissionWaiters(ctx context.Context, taskIDs []string, waiterIDs []string) error { if len(taskIDs) == 0 { return nil } if len(taskIDs) != len(waiterIDs) { return errors.New("task admission waiter renewal requires matching task and waiter IDs") } _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_admissions admission SET waiter_lease_expires_at = now() + $3::interval, updated_at = now() FROM unnest($1::uuid[], $2::text[]) AS owned(task_id, waiter_id) WHERE admission.task_id = owned.task_id AND admission.waiter_id = owned.waiter_id AND admission.mode = 'sync' AND admission.status = 'waiting'`, taskIDs, waiterIDs, admissionWaiterLeaseTTL.String(), ) return err } func (s *Store) DeleteTaskAdmission(ctx context.Context, taskID string) error { tx, err := s.pool.Begin(ctx) if err != nil { return err } defer tx.Rollback(ctx) if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil { return err } if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases SET released_at = now() 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 := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return err } return tx.Commit(ctx) } func (s *Store) ActiveTaskAdmissionLeases(ctx context.Context, taskID string) ([]ConcurrencyLease, error) { rows, err := s.pool.Query(ctx, ` SELECT id::text, GREATEST(EXTRACT(EPOCH FROM expires_at - acquired_at), 1)::float8 FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL AND expires_at > now() ORDER BY scope_type, scope_key, id`, taskID) if err != nil { return nil, err } defer rows.Close() leases := make([]ConcurrencyLease, 0) for rows.Next() { var lease ConcurrencyLease var ttlSeconds float64 if err := rows.Scan(&lease.ID, &ttlSeconds); err != nil { return nil, err } lease.TTL = time.Duration(ttlSeconds * float64(time.Second)) leases = append(leases, lease) } return leases, rows.Err() } func (s *Store) ListWaitingAsyncAdmissionTaskIDs(ctx context.Context, limit int) ([]string, error) { if limit <= 0 || limit > 1000 { limit = 100 } rows, err := s.pool.Query(ctx, ` SELECT task_id::text FROM gateway_task_admissions WHERE status = 'waiting' AND mode = 'async' ORDER BY priority ASC, enqueued_at ASC, task_id ASC LIMIT $1`, limit) if err != nil { return nil, err } defer rows.Close() taskIDs := make([]string, 0) for rows.Next() { var taskID string if err := rows.Scan(&taskID); err != nil { return nil, err } taskIDs = append(taskIDs, taskID) } return taskIDs, rows.Err() } // ListWaitingTaskAdmissionIDs returns a bounded set of FIFO leaders across // platform-model and user-group queues. It is used after capacity is released // so API processes wake only plausible queue heads instead of every // synchronous waiter. func (s *Store) ListWaitingTaskAdmissionIDs(ctx context.Context, limit int) ([]string, error) { if limit <= 0 { limit = 256 } rows, err := s.pool.Query(ctx, ` WITH ranked AS ( SELECT task_id, priority, enqueued_at, row_number() OVER ( PARTITION BY platform_model_id ORDER BY priority ASC, enqueued_at ASC, task_id ASC ) AS platform_rank, row_number() OVER ( PARTITION BY user_group_id ORDER BY priority ASC, enqueued_at ASC, task_id ASC ) AS group_rank FROM gateway_task_admissions WHERE status = 'waiting' AND mode = 'sync' ) SELECT task_id::text FROM ranked WHERE platform_rank <= $1 AND group_rank <= $1 ORDER BY priority ASC, enqueued_at ASC, task_id ASC LIMIT $1`, limit) if err != nil { return nil, err } defer rows.Close() taskIDs := make([]string, 0, limit) for rows.Next() { var taskID string if err := rows.Scan(&taskID); err != nil { return nil, err } taskIDs = append(taskIDs, taskID) } return taskIDs, rows.Err() } func (s *Store) GetTaskAdmission(ctx context.Context, taskID string) (TaskAdmission, error) { admission, found, err := loadTaskAdmission(ctx, s.pool, taskID) if err != nil { return TaskAdmission{}, err } if !found { return TaskAdmission{}, pgx.ErrNoRows } return admission, nil } func (s *Store) TaskAdmissionMetrics(ctx context.Context) (TaskAdmissionMetricsSnapshot, error) { var snapshot TaskAdmissionMetricsSnapshot err := s.pool.QueryRow(ctx, ` SELECT COUNT(*) FILTER (WHERE status = 'waiting')::int, COUNT(*) FILTER (WHERE status = 'waiting' AND mode = 'sync')::int, COUNT(*) FILTER (WHERE status = 'waiting' AND mode = 'async')::int, COALESCE(EXTRACT(EPOCH FROM now() - MIN(enqueued_at) FILTER (WHERE status = 'waiting')), 0)::float8, COUNT(*) FILTER ( WHERE status = 'waiting' AND mode = 'sync' AND waiter_lease_expires_at <= now() )::int, COUNT(*) FILTER ( WHERE status = 'waiting' AND wait_deadline_at <= now() )::int FROM gateway_task_admissions`).Scan( &snapshot.QueueDepth, &snapshot.WaitingSync, &snapshot.WaitingAsync, &snapshot.OldestWaitSeconds, &snapshot.ExpiredWaiterBacklog, &snapshot.ExpiredDeadlineBacklog, ) return snapshot, err } func (s *Store) ReapExpiredTaskAdmissions(ctx context.Context, limit int) (TaskAdmissionReapResult, error) { if limit <= 0 || limit > 1000 { limit = 500 } tx, err := s.pool.Begin(ctx) if err != nil { return TaskAdmissionReapResult{}, err } defer tx.Rollback(ctx) rows, err := tx.Query(ctx, ` SELECT task_id::text FROM gateway_task_admissions WHERE status = 'waiting' AND ( wait_deadline_at <= now() OR (mode = 'sync' AND waiter_lease_expires_at <= now()) ) ORDER BY wait_deadline_at ASC, task_id ASC LIMIT $1`, limit) if err != nil { return TaskAdmissionReapResult{}, err } expired := make([]string, 0) for rows.Next() { var taskID string if err := rows.Scan(&taskID); err != nil { rows.Close() return TaskAdmissionReapResult{}, err } expired = append(expired, taskID) } if err := rows.Err(); err != nil { rows.Close() return TaskAdmissionReapResult{}, err } rows.Close() result := TaskAdmissionReapResult{} reaped := 0 for _, taskID := range expired { if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil { return TaskAdmissionReapResult{}, err } var deadlineExpired bool err := tx.QueryRow(ctx, ` DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid AND status = 'waiting' AND ( wait_deadline_at <= now() OR (mode = 'sync' AND waiter_lease_expires_at <= now()) ) RETURNING wait_deadline_at <= now()`, taskID).Scan(&deadlineExpired) if errors.Is(err, pgx.ErrNoRows) { continue } if err != nil { return TaskAdmissionReapResult{}, err } code := "client_disconnected" message := "synchronous queue waiter disconnected" status := "cancelled" if deadlineExpired { code = "queue_timeout" message = ErrQueueTimeout.Error() status = "failed" result.ExpiredDeadlines++ } else { result.ExpiredWaiters++ } if _, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = $2, error = $3, error_code = $4, error_message = $3, finished_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'queued'`, taskID, status, message, code); err != nil { return TaskAdmissionReapResult{}, err } reaped++ } if reaped > 0 { if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil { return TaskAdmissionReapResult{}, err } } if err := tx.Commit(ctx); err != nil { return TaskAdmissionReapResult{}, err } return result, nil } // ListenTaskAdmissionNotifications holds one pooled PostgreSQL connection for // the process and forwards NOTIFY payloads as wake-up hints. Callers must still // retain periodic polling because LISTEN notifications are not durable. func (s *Store) ListenTaskAdmissionNotifications(ctx context.Context, onNotification func(string)) error { connection, err := s.pool.Acquire(ctx) if err != nil { return err } defer connection.Release() if _, err := connection.Exec(ctx, `LISTEN gateway_task_admission`); err != nil { return err } for { notification, err := connection.Conn().WaitForNotification(ctx) if err != nil { return err } if onNotification != nil { onNotification(notification.Payload) } } } func validateTaskAdmissionInput(input TaskAdmissionInput) error { if strings.TrimSpace(input.TaskID) == "" || strings.TrimSpace(input.PlatformID) == "" || strings.TrimSpace(input.PlatformModelID) == "" { return errors.New("task admission requires task, platform, and platform model") } if input.Mode != "sync" && input.Mode != "async" { return errors.New("task admission mode must be sync or async") } if input.Mode == "sync" && strings.TrimSpace(input.WaiterID) == "" { return errors.New("synchronous task admission requires a waiter ID") } return nil } func normalizedAdmissionScopes(scopes []AdmissionScope) []AdmissionScope { out := make([]AdmissionScope, 0, len(scopes)) seen := make(map[string]struct{}) for _, scope := range scopes { scope.ScopeType = strings.TrimSpace(scope.ScopeType) scope.ScopeKey = strings.TrimSpace(scope.ScopeKey) if scope.ScopeType == "" || scope.ScopeKey == "" { continue } if scope.Amount <= 0 { scope.Amount = 1 } if scope.LeaseTTLSeconds <= 0 { scope.LeaseTTLSeconds = 120 } key := admissionLockKey(scope) if _, ok := seen[key]; ok { continue } seen[key] = struct{}{} out = append(out, scope) } sort.Slice(out, func(i, j int) bool { return admissionLockKey(out[i]) < admissionLockKey(out[j]) }) return out } func admissionLockKey(scope AdmissionScope) string { return fmt.Sprintf("task-admission:%d:%s:%d:%s", len(scope.ScopeType), scope.ScopeType, len(scope.ScopeKey), scope.ScopeKey) } type admissionScopeState struct { Scope AdmissionScope Active float64 Waiting int Saturated bool NextLeaseExpiration time.Time } func admissionScopeStatesTx(ctx context.Context, tx pgx.Tx, scopes []AdmissionScope) ([]admissionScopeState, error) { states := make([]admissionScopeState, 0, len(scopes)) for _, scope := range scopes { state := admissionScopeState{Scope: scope} if scope.ConcurrentLimit > 0 { if err := tx.QueryRow(ctx, ` SELECT COALESCE(SUM(lease_value), 0)::float8, COALESCE(MIN(expires_at), now() + interval '1 second') FROM gateway_concurrency_leases WHERE scope_type = $1 AND scope_key = $2 AND released_at IS NULL AND expires_at > now()`, scope.ScopeType, scope.ScopeKey).Scan(&state.Active, &state.NextLeaseExpiration); err != nil { return nil, err } state.Saturated = state.Active+scope.Amount > scope.ConcurrentLimit } waiting, err := admissionScopeWaitingTx(ctx, tx, scope) if err != nil { return nil, err } state.Waiting = waiting states = append(states, state) } return states, nil } func admissionScopeWaitingTx(ctx context.Context, tx pgx.Tx, scope AdmissionScope) (int, error) { var waiting int switch scope.ScopeType { case "platform_model": err := tx.QueryRow(ctx, ` SELECT COUNT(*) FROM gateway_task_admissions WHERE status = 'waiting' AND platform_model_id::text = $1`, scope.ScopeKey).Scan(&waiting) return waiting, err case "user_group": err := tx.QueryRow(ctx, ` SELECT COUNT(*) FROM gateway_task_admissions WHERE status = 'waiting' AND user_group_id::text = $1`, scope.ScopeKey).Scan(&waiting) return waiting, err case "worker_capacity": err := tx.QueryRow(ctx, ` SELECT COUNT(*) FROM gateway_task_admissions WHERE status = 'waiting' AND mode = 'async'`).Scan(&waiting) return waiting, err default: return 0, nil } } func admissionDeadline(scopes []AdmissionScope) time.Time { maxWait := 0 for _, scope := range scopes { if scope.QueueLimit <= 0 { continue } wait := scope.MaxWaitSeconds if wait <= 0 { wait = DefaultQueueMaxWaitSeconds } if maxWait == 0 || wait < maxWait { maxWait = wait } } if maxWait == 0 { // A task may be registered before a River execution slot is available // even when business queueing is disabled. Keep that infrastructure wait // separate from the policy's default 600-second queue timeout. return time.Now().Add(executionSlotWaitMax) } return time.Now().Add(time.Duration(maxWait) * time.Second) } func concurrencyAdmissionError(state admissionScopeState) error { return &RateLimitExceededError{ ScopeType: state.Scope.ScopeType, ScopeKey: state.Scope.ScopeKey, ScopeName: state.Scope.ScopeName, ScopeMetadata: state.Scope.ScopeMetadata, Metric: "concurrent", Limit: state.Scope.ConcurrentLimit, Amount: state.Scope.Amount, Current: state.Active, Used: state.Active, Projected: state.Active + state.Scope.Amount, ResetAt: state.NextLeaseExpiration, Policy: state.Scope.Policy, Message: "concurrency limit is saturated and queueing is disabled", RetryAfter: queueRetryAfter(state.NextLeaseExpiration), Retryable: true, } } func queueFullAdmissionError(state admissionScopeState) error { return &RateLimitExceededError{ ScopeType: state.Scope.ScopeType, ScopeKey: state.Scope.ScopeKey, ScopeName: state.Scope.ScopeName, ScopeMetadata: state.Scope.ScopeMetadata, Metric: "queue_size", Reason: "queue_full", Limit: float64(state.Scope.QueueLimit), Amount: 1, Current: float64(state.Waiting), Used: float64(state.Waiting), Projected: float64(state.Waiting + 1), ResetAt: state.NextLeaseExpiration, Policy: state.Scope.Policy, Message: "task admission queue is full", RetryAfter: queueRetryAfter(state.NextLeaseExpiration), Retryable: true, QueueDepth: state.Waiting, QueueLimit: state.Scope.QueueLimit, } } func queueRetryAfter(next time.Time) time.Duration { if next.IsZero() { return time.Second } wait := time.Until(next) if wait < time.Second { return time.Second } if wait > 30*time.Second { return 30 * time.Second } return wait } func insertTaskAdmissionTx(ctx context.Context, tx pgx.Tx, input TaskAdmissionInput, deadline time.Time) (TaskAdmission, error) { waiterID := strings.TrimSpace(input.WaiterID) var waiterLease any if input.Mode == "sync" { waiterLease = time.Now().Add(admissionWaiterLeaseTTL) } row := tx.QueryRow(ctx, ` INSERT INTO gateway_task_admissions ( task_id, platform_id, platform_model_id, user_group_id, queue_key, mode, status, priority, wait_deadline_at, waiter_id, waiter_lease_expires_at ) VALUES ( $1::uuid, $2::uuid, $3::uuid, NULLIF($4, '')::uuid, $5, $6, 'waiting', $7, $8, NULLIF($9, ''), $10 ) ON CONFLICT (task_id) DO UPDATE SET waiter_id = CASE WHEN gateway_task_admissions.mode = 'sync' THEN EXCLUDED.waiter_id ELSE gateway_task_admissions.waiter_id END, waiter_lease_expires_at = CASE WHEN gateway_task_admissions.mode = 'sync' THEN EXCLUDED.waiter_lease_expires_at ELSE gateway_task_admissions.waiter_lease_expires_at END, updated_at = now() RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at`, input.TaskID, input.PlatformID, input.PlatformModelID, input.UserGroupID, input.QueueKey, input.Mode, input.Priority, deadline, waiterID, waiterLease) return scanTaskAdmission(row) } func renewTaskAdmissionWaiterTx(ctx context.Context, tx pgx.Tx, taskID string, waiterID string) (TaskAdmission, error) { row := tx.QueryRow(ctx, ` UPDATE gateway_task_admissions SET waiter_id = $2, waiter_lease_expires_at = now() + $3::interval, updated_at = now() WHERE task_id = $1::uuid AND mode = 'sync' AND status = 'waiting' RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at`, taskID, waiterID, admissionWaiterLeaseTTL.String()) return scanTaskAdmission(row) } func markTaskAdmissionAdmittedTx(ctx context.Context, tx pgx.Tx, taskID string) (TaskAdmission, error) { row := tx.QueryRow(ctx, ` UPDATE gateway_task_admissions SET status = 'admitted', admitted_at = now(), waiter_lease_expires_at = NULL, updated_at = now() WHERE task_id = $1::uuid AND status = 'waiting' RETURNING task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at`, taskID) return scanTaskAdmission(row) } func loadTaskAdmissionTx(ctx context.Context, tx pgx.Tx, taskID string) (TaskAdmission, bool, error) { return loadTaskAdmission(ctx, tx, taskID) } type admissionQuerier interface { QueryRow(context.Context, string, ...any) pgx.Row } func loadTaskAdmission(ctx context.Context, q admissionQuerier, taskID string) (TaskAdmission, bool, error) { row := q.QueryRow(ctx, ` SELECT task_id::text, platform_id::text, platform_model_id::text, COALESCE(user_group_id::text, ''), queue_key, mode, status, priority, enqueued_at, wait_deadline_at, admitted_at, COALESCE(waiter_id, ''), waiter_lease_expires_at FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID) admission, err := scanTaskAdmission(row) if errors.Is(err, pgx.ErrNoRows) { return TaskAdmission{}, false, nil } return admission, err == nil, err } func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) { var admission TaskAdmission var admittedAt sql.NullTime var waiterLeaseExpires sql.NullTime err := row.Scan( &admission.TaskID, &admission.PlatformID, &admission.PlatformModelID, &admission.UserGroupID, &admission.QueueKey, &admission.Mode, &admission.Status, &admission.Priority, &admission.EnqueuedAt, &admission.WaitDeadlineAt, &admittedAt, &admission.WaiterID, &waiterLeaseExpires, ) if admittedAt.Valid { admission.AdmittedAt = admittedAt.Time } if waiterLeaseExpires.Valid { admission.WaiterLeaseExpires = waiterLeaseExpires.Time } return admission, err } func isTaskAdmissionHeadTx(ctx context.Context, tx pgx.Tx, admission TaskAdmission, scopes []AdmissionScope) (bool, error) { for _, scope := range scopes { var head string switch scope.ScopeType { case "platform_model": if err := tx.QueryRow(ctx, ` SELECT task_id::text FROM gateway_task_admissions WHERE platform_model_id = $1::uuid AND status = 'waiting' ORDER BY priority ASC, enqueued_at ASC, task_id ASC LIMIT 1`, admission.PlatformModelID).Scan(&head); err != nil { return false, err } case "user_group": if admission.UserGroupID == "" { continue } if err := tx.QueryRow(ctx, ` SELECT task_id::text FROM gateway_task_admissions WHERE user_group_id = $1::uuid AND status = 'waiting' ORDER BY priority ASC, enqueued_at ASC, task_id ASC LIMIT 1`, admission.UserGroupID).Scan(&head); err != nil { return false, err } default: continue } if head != admission.TaskID { return false, nil } } return true, nil } func expireTaskAdmissionTx(ctx context.Context, tx pgx.Tx, taskID string) error { if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil { return err } _, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = 'failed', error = 'task admission queue wait timed out', error_code = 'queue_timeout', error_message = 'task admission queue wait timed out', finished_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'queued'`, taskID) return err } func notifyTaskAdmissionTx(ctx context.Context, tx pgx.Tx, taskID string) error { _, err := tx.Exec(ctx, `SELECT pg_notify($1, $2)`, admissionNotifyChannel, taskID) return err }