perf(worker): 持久化准入快照提升队列吞吐

原因:线上 P24 验收显示 Worker 在每轮准入前逐条恢复媒体、重算候选和查询 attempts,跨地域查询导致 48 个执行槽长期只能使用约 8 至 16 个。\n\n影响:任务首次排队时保存鉴权后的 admission scope 快照,Worker 单次批量读取并只刷新动态总容量;显式重新选路保留完整慢路径。验收模拟器改到非数据库专用 Worker 节点,运行期 85% 作为立即中止硬门禁,80% 继续作为容量认证目标。\n\n验证:Go 全量测试、go vet、govulncheck、真实 PostgreSQL 迁移与跨 Store 集成测试、ShellCheck、发布及验收脚本、OpenAPI、前端 lint/test/build、pnpm audit high、Compose 与 Kubernetes 渲染均通过。
This commit is contained in:
2026-08-01 21:25:54 +08:00
parent c89c56ca65
commit 56a5176c92
10 changed files with 274 additions and 100 deletions
+73 -59
View File
@@ -34,8 +34,8 @@ const (
asyncWorkerQueueLimit = 10000
asyncWorkerMaxWaitSeconds = 24 * 60 * 60
// The dispatcher scans enough FIFO entries to fill the certified 2 x P32
// capacity. Store admission commits each item independently so this scan
// size is never also a multi-task transaction size.
// capacity. Durable scope snapshots make this one bulk read; atomic commits
// remain bounded by the separately configured microbatch size.
asyncAdmissionBatchMax = 64
acceptanceQueueLimit = 10000
acceptanceQueueMaxWait = 15 * 60
@@ -690,48 +690,80 @@ func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) e
return nil
}
func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, tasks []store.GatewayTask) (bool, error) {
if len(tasks) == 0 {
func admissionInputFromSnapshot(admission store.TaskAdmission, workerCapacity int) store.TaskAdmissionInput {
scopes := append([]store.AdmissionScope(nil), admission.Scopes...)
for index := range scopes {
if scopes[index].ScopeType == "worker_capacity" && scopes[index].ScopeKey == asyncWorkerCapacityScopeKey {
scopes[index].ConcurrentLimit = float64(workerCapacity)
}
}
return store.TaskAdmissionInput{
TaskID: admission.TaskID,
PlatformID: admission.PlatformID,
PlatformModelID: admission.PlatformModelID,
UserGroupID: admission.UserGroupID,
QueueKey: admission.QueueKey,
Mode: admission.Mode,
Priority: admission.Priority,
Scopes: scopes,
}
}
func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []store.TaskAdmission) (bool, error) {
if len(admissions) == 0 {
return false, nil
}
inputs := make([]store.TaskAdmissionInput, 0, len(tasks))
tasksByID := make(map[string]store.GatewayTask, len(tasks))
for _, task := range tasks {
user := authUserFromTask(task)
current, err := s.store.GetTaskAdmission(ctx, task.ID)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
continue
}
return false, err
}
plan, err := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, user, &current)
if err != nil {
return false, err
}
if !plan.Eligible {
if err := s.EnqueueAsyncTask(ctx, task); err != nil {
return false, err
}
continue
}
input, err := s.prepareTaskAdmissionInput(ctx, task, plan, "")
if err != nil {
return false, err
}
inputs = append(inputs, input)
tasksByID[task.ID] = task
workerCapacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx)
if err != nil {
return false, err
}
if workerCapacity <= 0 {
return true, nil
}
saturated := false
microbatchSize := s.cfg.AsyncAdmissionMicrobatchSize
if microbatchSize <= 0 {
microbatchSize = 1
}
for start := 0; start < len(inputs) && !saturated; start += microbatchSize {
end := min(start+microbatchSize, len(inputs))
for start := 0; start < len(admissions); start += microbatchSize {
end := min(start+microbatchSize, len(admissions))
inputs := make([]store.TaskAdmissionInput, 0, end-start)
tasksByID := make(map[string]store.GatewayTask, end-start)
for _, admission := range admissions[start:end] {
if admission.ReselectRequestedAt.IsZero() && len(admission.Scopes) > 0 {
inputs = append(inputs, admissionInputFromSnapshot(admission, workerCapacity))
tasksByID[admission.TaskID] = store.GatewayTask{ID: admission.TaskID}
continue
}
task, loadErr := s.store.GetTask(ctx, admission.TaskID)
if loadErr != nil {
if errors.Is(loadErr, pgx.ErrNoRows) {
continue
}
return false, loadErr
}
plan, planErr := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, authUserFromTask(task), &admission)
if planErr != nil {
return false, planErr
}
if !plan.Eligible {
if enqueueErr := s.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
return false, enqueueErr
}
continue
}
input := taskAdmissionInput(task, plan, "")
if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil && !errors.Is(rebindErr, pgx.ErrNoRows) {
return false, rebindErr
}
inputs = append(inputs, input)
tasksByID[task.ID] = task
}
if len(inputs) == 0 {
continue
}
outcomes, err := s.store.TryTaskAdmissionAtomicBatchWithAdmittedHook(
ctx,
inputs[start:end],
inputs,
func(tx pgx.Tx, input store.TaskAdmissionInput) error {
task, ok := tasksByID[input.TaskID]
if !ok {
@@ -753,12 +785,11 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, tasks []store.G
return false, outcome.Err
}
if !outcome.Result.Admitted {
saturated = true
break
return true, nil
}
}
}
return saturated, nil
return false, nil
}
func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskID string) bool {
@@ -789,34 +820,17 @@ func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
}
for {
batchLimit := asyncAdmissionBatchLimit(s.cfg.AsyncWorkerHardLimit)
taskIDs, err := s.store.ListWaitingAsyncAdmissionTaskIDs(ctx, batchLimit)
admissions, err := s.store.ListWaitingAsyncAdmissions(ctx, batchLimit)
if err != nil {
s.logger.Warn("list waiting async admissions failed", "error", err)
break
}
if len(taskIDs) == 0 {
if len(admissions) == 0 {
break
}
tasks := make([]store.GatewayTask, 0, len(taskIDs))
for _, taskID := range taskIDs {
task, err := s.store.GetTask(ctx, taskID)
if err != nil {
if !errors.Is(err, pgx.ErrNoRows) {
s.logger.Warn("load waiting async task failed", "taskID", taskID, "error", err)
}
continue
}
if task.Status != "queued" {
continue
}
tasks = append(tasks, task)
}
if len(tasks) == 0 {
break
}
saturated, err := s.dispatchWaitingAsyncTasks(ctx, tasks)
saturated, err := s.dispatchWaitingAsyncTasks(ctx, admissions)
if err != nil {
s.logger.Warn("dispatch waiting async admission batch failed", "tasks", len(tasks), "error", err)
s.logger.Warn("dispatch waiting async admission batch failed", "tasks", len(admissions), "error", err)
break
}
if saturated {
@@ -2,10 +2,43 @@ package runner
import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAdmissionInputFromSnapshotRefreshesWorkerCapacity(t *testing.T) {
admission := store.TaskAdmission{
TaskID: "task-1",
PlatformID: "platform-1",
PlatformModelID: "model-1",
UserGroupID: "group-1",
QueueKey: "queue-1",
Mode: "async",
Priority: 7,
Scopes: []store.AdmissionScope{
{ScopeType: "platform_model", ScopeKey: "model-1", ConcurrentLimit: 10},
{ScopeType: "worker_capacity", ScopeKey: "global", ConcurrentLimit: 48},
},
ReselectRequestedAt: time.Time{},
}
input := admissionInputFromSnapshot(admission, 24)
if input.TaskID != admission.TaskID || input.PlatformModelID != admission.PlatformModelID || len(input.Scopes) != 2 {
t.Fatalf("unexpected admission input: %+v", input)
}
if input.Scopes[1].ConcurrentLimit != 24 {
t.Fatalf("worker capacity=%v, want 24", input.Scopes[1].ConcurrentLimit)
}
if input.Scopes[0].ConcurrentLimit != 10 {
t.Fatalf("business scope changed: %+v", input.Scopes[0])
}
if admission.Scopes[1].ConcurrentLimit != 48 {
t.Fatalf("snapshot was mutated: %+v", admission.Scopes[1])
}
}
func TestDistributedAdmissionModelTypeBoundary(t *testing.T) {
for _, modelType := range []string{
"image_generate", "image_edit", "image_analysis", "image_vectorize",
+119 -22
View File
@@ -3,6 +3,7 @@ package store
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"math"
@@ -64,19 +65,21 @@ type TaskAdmissionInput struct {
}
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
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
Scopes []AdmissionScope
ReselectRequestedAt time.Time
}
type TaskAdmissionResult struct {
@@ -697,6 +700,10 @@ func (s *Store) rebindWaitingTaskAdmissionOnce(ctx context.Context, input TaskAd
return TaskAdmission{}, admissionErr
}
deadline := admissionDeadline(scopes)
scopeSnapshot, err := json.Marshal(scopes)
if err != nil {
return TaskAdmission{}, fmt.Errorf("marshal rebound task admission scopes: %w", err)
}
row := tx.QueryRow(ctx, `
UPDATE gateway_task_admissions
SET platform_id = $2::uuid,
@@ -704,6 +711,7 @@ SET platform_id = $2::uuid,
user_group_id = NULLIF($4, '')::uuid,
queue_key = $5,
wait_deadline_at = LEAST(wait_deadline_at, $6),
scope_snapshot = $7::jsonb,
reselect_requested_at = NULL,
updated_at = now()
WHERE task_id = $1::uuid
@@ -711,13 +719,15 @@ 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`,
COALESCE(waiter_id, ''), waiter_lease_expires_at,
scope_snapshot, reselect_requested_at`,
input.TaskID,
input.PlatformID,
input.PlatformModelID,
input.UserGroupID,
input.QueueKey,
deadline,
scopeSnapshot,
)
migrated, err := scanTaskAdmission(row)
if err != nil {
@@ -753,6 +763,10 @@ WHERE task_id = $1::uuid
if input.Mode == "sync" {
waiterLease = time.Now().Add(admissionWaiterLeaseTTL)
}
scopeSnapshot, err := json.Marshal(normalizedAdmissionScopes(input.Scopes))
if err != nil {
return TaskAdmission{}, fmt.Errorf("marshal reset task admission scopes: %w", err)
}
row := tx.QueryRow(ctx, `
UPDATE gateway_task_admissions
SET status = 'waiting',
@@ -764,12 +778,14 @@ SET status = 'waiting',
wait_deadline_at = LEAST(wait_deadline_at, $6),
waiter_id = NULLIF($7, ''),
waiter_lease_expires_at = $8,
scope_snapshot = $9::jsonb,
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`,
COALESCE(waiter_id, ''), waiter_lease_expires_at,
scope_snapshot, reselect_requested_at`,
input.TaskID,
input.PlatformID,
input.PlatformModelID,
@@ -778,6 +794,7 @@ RETURNING task_id::text, platform_id::text, platform_model_id::text,
admissionDeadline(input.Scopes),
strings.TrimSpace(input.WaiterID),
waiterLease,
scopeSnapshot,
)
return scanTaskAdmission(row)
}
@@ -933,6 +950,62 @@ LIMIT $1`, limit)
return taskIDs, rows.Err()
}
// ListWaitingAsyncAdmissions returns the durable routing and admission-scope
// snapshots needed by the Worker dispatcher in one round trip. Normal queued
// tasks must not be fully rehydrated and routed again merely to reserve an
// execution slot; only rows explicitly marked for reselection take that slower
// path in the runner.
func (s *Store) ListWaitingAsyncAdmissions(ctx context.Context, limit int) ([]TaskAdmission, error) {
if limit <= 0 || limit > 1000 {
limit = 100
}
rows, err := s.pool.Query(ctx, `
SELECT admission.task_id::text, admission.platform_id::text, admission.platform_model_id::text,
COALESCE(admission.user_group_id::text, ''), admission.queue_key, admission.mode,
admission.status, admission.priority, admission.enqueued_at, admission.wait_deadline_at,
admission.admitted_at, COALESCE(admission.waiter_id, ''), admission.waiter_lease_expires_at,
admission.scope_snapshot, admission.reselect_requested_at
FROM gateway_task_admissions admission
JOIN gateway_tasks task ON task.id = admission.task_id
LEFT JOIN river_job job ON job.id = task.river_job_id
WHERE admission.mode = 'async'
AND task.status = 'queued'
AND task.next_run_at <= now()
AND (
admission.status = 'waiting'
OR (
admission.status = 'admitted'
AND NOT EXISTS (
SELECT 1
FROM gateway_concurrency_leases lease
WHERE lease.task_id = admission.task_id
AND lease.released_at IS NULL
AND lease.expires_at > now()
)
)
)
AND (
task.river_job_id IS NULL
OR job.id IS NULL
OR job.state <> 'running'
)
ORDER BY admission.priority ASC, admission.enqueued_at ASC, admission.task_id ASC
LIMIT $1`, limit)
if err != nil {
return nil, err
}
defer rows.Close()
admissions := make([]TaskAdmission, 0, limit)
for rows.Next() {
admission, scanErr := scanTaskAdmission(rows)
if scanErr != nil {
return nil, scanErr
}
admissions = append(admissions, admission)
}
return admissions, rows.Err()
}
// ListWaitingTaskAdmissionIDs returns one FIFO leader for every independent
// platform-model queue, additionally requiring the task to lead its user-group
// queue when one exists. It is used after capacity is released so API
@@ -1323,6 +1396,10 @@ func queueRetryAfter(next time.Time) time.Duration {
func insertTaskAdmissionTx(ctx context.Context, tx pgx.Tx, input TaskAdmissionInput, deadline time.Time) (TaskAdmission, error) {
waiterID := strings.TrimSpace(input.WaiterID)
scopeSnapshot, err := json.Marshal(normalizedAdmissionScopes(input.Scopes))
if err != nil {
return TaskAdmission{}, fmt.Errorf("marshal task admission scopes: %w", err)
}
var waiterLease any
if input.Mode == "sync" {
waiterLease = time.Now().Add(admissionWaiterLeaseTTL)
@@ -1330,22 +1407,27 @@ func insertTaskAdmissionTx(ctx context.Context, tx pgx.Tx, input TaskAdmissionIn
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
priority, wait_deadline_at, waiter_id, waiter_lease_expires_at, scope_snapshot
)
VALUES (
$1::uuid, $2::uuid, $3::uuid, NULLIF($4, '')::uuid, $5, $6, 'waiting',
$7, $8, NULLIF($9, ''), $10
$7, $8, NULLIF($9, ''), $10, $11::jsonb
)
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,
scope_snapshot = CASE
WHEN gateway_task_admissions.scope_snapshot IS NULL OR gateway_task_admissions.scope_snapshot = '[]'::jsonb THEN EXCLUDED.scope_snapshot
ELSE gateway_task_admissions.scope_snapshot
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`,
COALESCE(waiter_id, ''), waiter_lease_expires_at,
scope_snapshot, reselect_requested_at`,
input.TaskID, input.PlatformID, input.PlatformModelID, input.UserGroupID,
input.QueueKey, input.Mode, input.Priority, deadline, waiterID, waiterLease)
input.QueueKey, input.Mode, input.Priority, deadline, waiterID, waiterLease, scopeSnapshot)
return scanTaskAdmission(row)
}
@@ -1361,7 +1443,8 @@ 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`,
COALESCE(waiter_id, ''), waiter_lease_expires_at,
scope_snapshot, reselect_requested_at`,
taskID, waiterID, admissionWaiterLeaseTTL.String())
return scanTaskAdmission(row)
}
@@ -1377,7 +1460,8 @@ 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`,
COALESCE(waiter_id, ''), waiter_lease_expires_at,
scope_snapshot, reselect_requested_at`,
taskID)
return scanTaskAdmission(row)
}
@@ -1395,7 +1479,8 @@ func loadTaskAdmission(ctx context.Context, q admissionQuerier, taskID string) (
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
COALESCE(waiter_id, ''), waiter_lease_expires_at,
scope_snapshot, reselect_requested_at
FROM gateway_task_admissions
WHERE task_id = $1::uuid`, taskID)
admission, err := scanTaskAdmission(row)
@@ -1409,6 +1494,8 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) {
var admission TaskAdmission
var admittedAt sql.NullTime
var waiterLeaseExpires sql.NullTime
var reselectRequestedAt sql.NullTime
var scopeSnapshot []byte
err := row.Scan(
&admission.TaskID,
&admission.PlatformID,
@@ -1423,6 +1510,8 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) {
&admittedAt,
&admission.WaiterID,
&waiterLeaseExpires,
&scopeSnapshot,
&reselectRequestedAt,
)
if admittedAt.Valid {
admission.AdmittedAt = admittedAt.Time
@@ -1430,6 +1519,14 @@ func scanTaskAdmission(row pgx.Row) (TaskAdmission, error) {
if waiterLeaseExpires.Valid {
admission.WaiterLeaseExpires = waiterLeaseExpires.Time
}
if reselectRequestedAt.Valid {
admission.ReselectRequestedAt = reselectRequestedAt.Time
}
if err == nil && len(scopeSnapshot) > 0 {
if unmarshalErr := json.Unmarshal(scopeSnapshot, &admission.Scopes); unmarshalErr != nil {
return TaskAdmission{}, fmt.Errorf("unmarshal task admission scopes: %w", unmarshalErr)
}
}
return admission, err
}
@@ -434,6 +434,23 @@ WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID)
if err != nil || queuedAdmission.Status != "waiting" {
t.Fatalf("atomic queued admission=%+v err=%v, want waiting", queuedAdmission, err)
}
if len(queuedAdmission.Scopes) == 0 {
t.Fatal("queued admission did not persist its scope snapshot")
}
waitingAdmissionSnapshots, err := first.ListWaitingAsyncAdmissions(ctx, 1000)
if err != nil {
t.Fatalf("list waiting async admission snapshots: %v", err)
}
var listedSnapshot *TaskAdmission
for index := range waitingAdmissionSnapshots {
if waitingAdmissionSnapshots[index].TaskID == queuedAtomicTask.ID {
listedSnapshot = &waitingAdmissionSnapshots[index]
break
}
}
if listedSnapshot == nil || len(listedSnapshot.Scopes) != len(queuedAdmission.Scopes) {
t.Fatalf("listed admission snapshot=%+v, want %d scopes", listedSnapshot, len(queuedAdmission.Scopes))
}
var riverJobID int64
if err := first.pool.QueryRow(ctx, `
SELECT
@@ -0,0 +1,9 @@
ALTER TABLE gateway_task_admissions
ADD COLUMN IF NOT EXISTS scope_snapshot jsonb;
ALTER TABLE gateway_task_admissions
ADD CONSTRAINT gateway_task_admissions_scope_snapshot_array_check
CHECK (scope_snapshot IS NULL OR jsonb_typeof(scope_snapshot) = 'array') NOT VALID;
ALTER TABLE gateway_task_admissions
VALIDATE CONSTRAINT gateway_task_admissions_scope_snapshot_array_check;