wangbo
9e4fc7362d
使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。 补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。 已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
1045 lines
31 KiB
Go
1045 lines
31 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
const (
|
|
admissionWaiterLeaseTTL = 15 * time.Second
|
|
admissionNotifyChannel = "gateway_task_admission"
|
|
)
|
|
|
|
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)
|
|
}
|
|
|
|
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, input.TaskID); 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
|
|
}
|
|
if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); 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
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
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, taskID); 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()
|
|
}
|
|
|
|
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
|
|
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 {
|
|
maxWait = DefaultQueueMaxWaitSeconds
|
|
}
|
|
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
|
|
}
|