easyai-ai-gateway/apps/api/internal/store/admission_queue.go

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
}