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:
@@ -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, ¤t)
|
||||
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",
|
||||
|
||||
@@ -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;
|
||||
Reference in New Issue
Block a user