BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
原因：线上 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:
wangbo
2026-08-01 21:25:54 +08:00
parent c89c56ca65
commit 56a5176c92
10 changed files with 274 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/admission_queue.go
+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
}
apps/api/internal/store/admission_queue_integration_test.go
+17
View File
@@ -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