perf(queue): 限制大媒体任务内存并消除准入惊群

将同步与异步非文本任务的准入唤醒改为按任务和全局队首推进，批量续租等待者，避免大量请求同时争抢 advisory lock 和数据库连接。对 Base64 请求解析、素材解码、上游媒体执行和结果物化增加分层并发限制，复用请求素材并释放重复 Gemini wire 数据；生产 API 默认入口解析并发 16，媒体物化并发 8。新增 1000 个同步 Gemini 图像编辑请求的模拟上游压力验收，10 MiB 输入和输出下全部成功，Heap 峰值增长约 2.58 GiB，并验证共享上传哈希、单次 attempt 与本地零落盘。验证：go test ./... -count=1；go vet ./...；gofmt -l 无输出；kubectl kustomize deploy/kubernetes/production；10 MiB Gemini Base64 千任务压力测试通过。
2026-07-30 13:13:32 +08:00
parent fd3b6bf042
commit ea58d21d03
26 changed files with 1980 additions and 74 deletions
@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"math"
 	"sort"
 	"strings"
 	"time"
@@ -15,6 +16,7 @@ import (
 const (
 	admissionWaiterLeaseTTL = 15 * time.Second
 	admissionNotifyChannel  = "gateway_task_admission"
+	executionSlotWaitMax    = 24 * time.Hour
 )

 var ErrQueueTimeout = errors.New("task admission queue wait timed out")
@@ -113,6 +115,111 @@ func (s *Store) TryTaskAdmissionWithAdmittedHook(
 	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,
@@ -196,7 +303,7 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
 		if err := expireTaskAdmissionTx(ctx, tx, input.TaskID); err != nil {
 			return TaskAdmissionResult{}, err
 		}
-		if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil {
+		if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil {
 			return TaskAdmissionResult{}, err
 		}
 		if err := tx.Commit(ctx); err != nil {
@@ -276,7 +383,10 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil {
 	if err != nil {
 		return TaskAdmissionResult{}, err
 	}
-	if err := notifyTaskAdmissionTx(ctx, tx, input.TaskID); err != nil {
+	// 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 {
@@ -312,6 +422,29 @@ func validateNewAdmissionCapacity(states []admissionScopeState) error {
 	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.
@@ -469,6 +602,32 @@ WHERE task_id = $1::uuid
 	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 {
@@ -488,7 +647,7 @@ WHERE task_id = $1::uuid
 	if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
 		return err
 	}
-	if err := notifyTaskAdmissionTx(ctx, tx, taskID); err != nil {
+	if err := notifyTaskAdmissionTx(ctx, tx, "*"); err != nil {
 		return err
 	}
 	return tx.Commit(ctx)
@@ -546,6 +705,52 @@ LIMIT $1`, limit)
 	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 {
@@ -793,6 +998,12 @@ 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
 	}
@@ -813,7 +1024,10 @@ func admissionDeadline(scopes []AdmissionScope) time.Time {
 		}
 	}
 	if maxWait == 0 {
-		maxWait = DefaultQueueMaxWaitSeconds
+		// 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)
 }