wangbo
6bab0f0749
为 PostgreSQL 连接增加可配置的事务空闲与锁等待超时，并在请求取消或 Worker 退出后使用独立有界上下文回滚事务。\n\n恢复过期任务时按批次使用 SKIP LOCKED，周期重建缺失或终态 River Job；锁等待超时只在确认尚未提交上游时安全重排，避免重复执行和重复结算。\n\n验证：完整 Go 测试、go vet、生产 Kustomize 渲染、gofmt 与 git diff --check 均通过。
796 lines
29 KiB
Go
796 lines
29 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
func TestDistributedTaskAdmissionQueueAcrossStores(t *testing.T) {
|
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
|
if databaseURL == "" {
|
|
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run task admission PostgreSQL integration tests")
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
|
defer cancel()
|
|
applyOIDCJITTestMigrations(t, ctx, databaseURL)
|
|
|
|
first, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect first store: %v", err)
|
|
}
|
|
defer first.Close()
|
|
second, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect second store: %v", err)
|
|
}
|
|
defer second.Close()
|
|
var databaseName string
|
|
if err := first.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
|
|
t.Fatalf("read test database name: %v", err)
|
|
}
|
|
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
|
t.Fatalf("refusing to use non-test database %q", databaseName)
|
|
}
|
|
|
|
suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
|
|
platform, err := first.CreatePlatform(ctx, CreatePlatformInput{
|
|
Provider: "queue-test",
|
|
PlatformKey: "queue-test-" + suffix,
|
|
Name: "Queue Test " + suffix,
|
|
AuthType: "none",
|
|
Status: "enabled",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("create platform: %v", err)
|
|
}
|
|
var platformModelID string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
INSERT INTO platform_models (
|
|
platform_id, model_name, provider_model_name, model_alias, model_type,
|
|
display_name, pricing_mode, enabled
|
|
)
|
|
VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
|
|
RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID); err != nil {
|
|
t.Fatalf("create platform model: %v", err)
|
|
}
|
|
taskIDs := make([]string, 0, 4)
|
|
t.Cleanup(func() {
|
|
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cleanupCancel()
|
|
for _, taskID := range taskIDs {
|
|
_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
|
|
}
|
|
_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = $1::uuid`, platformModelID)
|
|
_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
|
|
_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id LIKE $1`, "queue-test-"+suffix+"%")
|
|
})
|
|
|
|
createTask := func(async bool) GatewayTask {
|
|
t.Helper()
|
|
task, createErr := first.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generate",
|
|
Model: "queue-model-" + suffix,
|
|
RunMode: "production",
|
|
Async: async,
|
|
Request: map[string]any{"prompt": "queue admission test"},
|
|
}, &auth.User{ID: "queue-test-user-" + suffix, Source: "gateway"})
|
|
if createErr != nil {
|
|
t.Fatalf("create task: %v", createErr)
|
|
}
|
|
taskIDs = append(taskIDs, task.ID)
|
|
return task
|
|
}
|
|
scope := AdmissionScope{
|
|
ScopeType: "platform_model",
|
|
ScopeKey: platformModelID,
|
|
ScopeName: "queue-model",
|
|
ConcurrentLimit: 1,
|
|
Amount: 1,
|
|
LeaseTTLSeconds: 137,
|
|
QueueLimit: 2,
|
|
MaxWaitSeconds: 600,
|
|
}
|
|
inputFor := func(task GatewayTask, priority int, waiterID string) TaskAdmissionInput {
|
|
mode := "sync"
|
|
if task.AsyncMode {
|
|
mode = "async"
|
|
}
|
|
return TaskAdmissionInput{
|
|
TaskID: task.ID,
|
|
PlatformID: platform.ID,
|
|
PlatformModelID: platformModelID,
|
|
QueueKey: "queue-test:" + suffix,
|
|
Mode: mode,
|
|
Priority: priority,
|
|
WaiterID: waiterID,
|
|
Scopes: []AdmissionScope{scope},
|
|
}
|
|
}
|
|
|
|
running := createTask(false)
|
|
result, err := first.TryTaskAdmission(ctx, inputFor(running, 100, "waiter-running"))
|
|
if err != nil || !result.Admitted || len(result.Leases) != 1 {
|
|
t.Fatalf("first task admission = %+v, err=%v", result, err)
|
|
}
|
|
reloadedLeases, err := second.ActiveTaskAdmissionLeases(ctx, running.ID)
|
|
if err != nil || len(reloadedLeases) != 1 || reloadedLeases[0].TTL != 137*time.Second {
|
|
t.Fatalf("reloaded admission leases = %+v, err=%v, want one 137s lease", reloadedLeases, err)
|
|
}
|
|
lowerPriority := createTask(false)
|
|
result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
|
|
if err != nil || result.Admitted {
|
|
t.Fatalf("second task should wait: result=%+v err=%v", result, err)
|
|
}
|
|
higherPriorityAsync := createTask(true)
|
|
result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
|
|
if err != nil || result.Admitted {
|
|
t.Fatalf("third task should wait: result=%+v err=%v", result, err)
|
|
}
|
|
rejected := createTask(true)
|
|
_, err = second.TryTaskAdmission(ctx, inputFor(rejected, 30, ""))
|
|
var limitErr *RateLimitExceededError
|
|
if !errors.As(err, &limitErr) || limitErr.Metric != "queue_size" || limitErr.Reason != "queue_full" ||
|
|
limitErr.QueueDepth != 2 || limitErr.QueueLimit != 2 {
|
|
t.Fatalf("fourth task error = %#v, want queue_full depth 2/2", err)
|
|
}
|
|
|
|
var queueCounters int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT COUNT(*)
|
|
FROM gateway_rate_limit_counters
|
|
WHERE metric = 'queue_size'
|
|
AND scope_key = $1`, platformModelID).Scan(&queueCounters); err != nil {
|
|
t.Fatalf("count queue counters: %v", err)
|
|
}
|
|
if queueCounters != 0 {
|
|
t.Fatalf("queue_size wrote %d fixed-window counters", queueCounters)
|
|
}
|
|
var attempts int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT COUNT(*)
|
|
FROM gateway_task_attempts
|
|
WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil {
|
|
t.Fatalf("count attempts: %v", err)
|
|
}
|
|
if attempts != 0 {
|
|
t.Fatalf("waiting admissions created %d task attempts", attempts)
|
|
}
|
|
|
|
if err := first.DeleteTaskAdmission(ctx, running.ID); err != nil {
|
|
t.Fatalf("release running admission: %v", err)
|
|
}
|
|
result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
|
|
if err != nil {
|
|
t.Fatalf("try lower-priority task after release: %v", err)
|
|
}
|
|
if result.Admitted {
|
|
t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter")
|
|
}
|
|
result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
|
|
if err != nil || !result.Admitted {
|
|
t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err)
|
|
}
|
|
for _, taskID := range taskIDs {
|
|
_ = first.DeleteTaskAdmission(ctx, taskID)
|
|
}
|
|
|
|
concurrentTasks := []GatewayTask{createTask(false), createTask(true), createTask(false), createTask(true)}
|
|
type concurrentResult struct {
|
|
result TaskAdmissionResult
|
|
err error
|
|
}
|
|
start := make(chan struct{})
|
|
results := make(chan concurrentResult, len(concurrentTasks))
|
|
var waitGroup sync.WaitGroup
|
|
for index, task := range concurrentTasks {
|
|
waitGroup.Add(1)
|
|
go func(index int, task GatewayTask) {
|
|
defer waitGroup.Done()
|
|
<-start
|
|
target := first
|
|
if index%2 == 1 {
|
|
target = second
|
|
}
|
|
waiterID := ""
|
|
if !task.AsyncMode {
|
|
waiterID = "concurrent-waiter-" + task.ID
|
|
}
|
|
admission, admissionErr := target.TryTaskAdmission(ctx, inputFor(task, 100, waiterID))
|
|
results <- concurrentResult{result: admission, err: admissionErr}
|
|
}(index, task)
|
|
}
|
|
close(start)
|
|
waitGroup.Wait()
|
|
close(results)
|
|
admittedCount := 0
|
|
waitingCount := 0
|
|
queueFullCount := 0
|
|
for item := range results {
|
|
if item.err == nil && item.result.Admitted {
|
|
admittedCount++
|
|
continue
|
|
}
|
|
if item.err == nil {
|
|
waitingCount++
|
|
continue
|
|
}
|
|
if errors.As(item.err, &limitErr) && limitErr.Metric == "queue_size" {
|
|
queueFullCount++
|
|
continue
|
|
}
|
|
t.Fatalf("unexpected concurrent admission result: %+v err=%v", item.result, item.err)
|
|
}
|
|
if admittedCount != 1 || waitingCount != 2 || queueFullCount != 1 {
|
|
t.Fatalf("concurrent outcomes admitted=%d waiting=%d queue_full=%d, want 1/2/1", admittedCount, waitingCount, queueFullCount)
|
|
}
|
|
var activeLeases, waitingAdmissions int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT COUNT(*)
|
|
FROM gateway_concurrency_leases
|
|
WHERE scope_type = 'platform_model'
|
|
AND scope_key = $1
|
|
AND released_at IS NULL
|
|
AND expires_at > now()`, platformModelID).Scan(&activeLeases); err != nil {
|
|
t.Fatalf("count active leases: %v", err)
|
|
}
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT COUNT(*)
|
|
FROM gateway_task_admissions
|
|
WHERE platform_model_id = $1::uuid
|
|
AND status = 'waiting'`, platformModelID).Scan(&waitingAdmissions); err != nil {
|
|
t.Fatalf("count waiting admissions: %v", err)
|
|
}
|
|
if activeLeases != 1 || waitingAdmissions != 2 {
|
|
t.Fatalf("database bounds active=%d waiting=%d, want 1/2", activeLeases, waitingAdmissions)
|
|
}
|
|
|
|
var deadlineTaskID string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT task_id::text
|
|
FROM gateway_task_admissions
|
|
WHERE platform_model_id = $1::uuid
|
|
AND status = 'waiting'
|
|
ORDER BY task_id
|
|
LIMIT 1`, platformModelID).Scan(&deadlineTaskID); err != nil {
|
|
t.Fatalf("select deadline waiter: %v", err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_task_admissions
|
|
SET enqueued_at = now() - interval '2 seconds',
|
|
wait_deadline_at = now() - interval '1 second'
|
|
WHERE task_id = $1::uuid`, deadlineTaskID); err != nil {
|
|
t.Fatalf("expire queue deadline: %v", err)
|
|
}
|
|
reaped, err := second.ReapExpiredTaskAdmissions(ctx, 10)
|
|
if err != nil || reaped.ExpiredDeadlines < 1 {
|
|
t.Fatalf("deadline reap = %+v, err=%v", reaped, err)
|
|
}
|
|
var taskStatus, taskErrorCode string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT status, COALESCE(error_code, '')
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, deadlineTaskID).Scan(&taskStatus, &taskErrorCode); err != nil {
|
|
t.Fatalf("read deadline task: %v", err)
|
|
}
|
|
if taskStatus != "failed" || taskErrorCode != "queue_timeout" {
|
|
t.Fatalf("deadline task status=%s code=%s, want failed/queue_timeout", taskStatus, taskErrorCode)
|
|
}
|
|
|
|
for _, taskID := range taskIDs {
|
|
_ = first.DeleteTaskAdmission(ctx, taskID)
|
|
}
|
|
leaseHolder := createTask(false)
|
|
if result, err := first.TryTaskAdmission(ctx, inputFor(leaseHolder, 100, "lease-holder")); err != nil || !result.Admitted {
|
|
t.Fatalf("admit lease holder: result=%+v err=%v", result, err)
|
|
}
|
|
disconnectedWaiter := createTask(false)
|
|
if result, err := second.TryTaskAdmission(ctx, inputFor(disconnectedWaiter, 100, "disconnected-waiter")); err != nil || result.Admitted {
|
|
t.Fatalf("queue disconnected waiter: result=%+v err=%v", result, err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_task_admissions
|
|
SET waiter_lease_expires_at = now() - interval '1 second'
|
|
WHERE task_id = $1::uuid`, disconnectedWaiter.ID); err != nil {
|
|
t.Fatalf("expire waiter lease: %v", err)
|
|
}
|
|
reaped, err = second.ReapExpiredTaskAdmissions(ctx, 10)
|
|
if err != nil || reaped.ExpiredWaiters < 1 {
|
|
t.Fatalf("waiter reap = %+v, err=%v", reaped, err)
|
|
}
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT status, COALESCE(error_code, '')
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, disconnectedWaiter.ID).Scan(&taskStatus, &taskErrorCode); err != nil {
|
|
t.Fatalf("read disconnected task: %v", err)
|
|
}
|
|
if taskStatus != "cancelled" || taskErrorCode != "client_disconnected" {
|
|
t.Fatalf("disconnected task status=%s code=%s, want cancelled/client_disconnected", taskStatus, taskErrorCode)
|
|
}
|
|
if err := first.DeleteTaskAdmission(ctx, leaseHolder.ID); err != nil {
|
|
t.Fatalf("release lease holder admission: %v", err)
|
|
}
|
|
|
|
preSubmission := createTask(false)
|
|
claimed, err := first.ClaimTaskPreparation(ctx, preSubmission.ID, uuid.NewString(), time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("claim pre-submission task: %v", err)
|
|
}
|
|
result, err = first.TryTaskAdmission(ctx, inputFor(claimed, 100, "pre-submission"))
|
|
if err != nil || !result.Admitted {
|
|
t.Fatalf("admit pre-submission task: result=%+v err=%v", result, err)
|
|
}
|
|
var attemptID string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
INSERT INTO gateway_task_attempts (
|
|
task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
|
|
upstream_submission_status
|
|
)
|
|
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'not_submitted')
|
|
RETURNING id::text`,
|
|
preSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
|
|
).Scan(&attemptID); err != nil {
|
|
t.Fatalf("create pre-submission attempt: %v", err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
INSERT INTO gateway_task_param_preprocessing_logs (
|
|
task_id, attempt_id, attempt_no, model_type
|
|
)
|
|
VALUES ($1::uuid, $2::uuid, 1, 'image_generate')`, preSubmission.ID, attemptID); err != nil {
|
|
t.Fatalf("create pre-submission log: %v", err)
|
|
}
|
|
cancelled, changed, err := first.CancelTaskBeforeUpstreamSubmission(
|
|
ctx,
|
|
preSubmission.ID,
|
|
claimed.ExecutionToken,
|
|
"client disconnected before upstream submission",
|
|
)
|
|
if err != nil || !changed || cancelled.Status != "cancelled" || cancelled.ErrorCode != "client_disconnected" {
|
|
t.Fatalf("cancel pre-submission task = %+v changed=%v err=%v", cancelled, changed, err)
|
|
}
|
|
var remainingAdmissions, remainingLeases, remainingAttempts, remainingPreprocessing int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
|
|
(SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL),
|
|
(SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = $1::uuid),
|
|
(SELECT COUNT(*) FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid)`,
|
|
preSubmission.ID,
|
|
).Scan(&remainingAdmissions, &remainingLeases, &remainingAttempts, &remainingPreprocessing); err != nil {
|
|
t.Fatalf("read pre-submission cleanup: %v", err)
|
|
}
|
|
if remainingAdmissions != 0 || remainingLeases != 0 || remainingAttempts != 0 || remainingPreprocessing != 0 {
|
|
t.Fatalf(
|
|
"pre-submission cleanup admissions=%d leases=%d attempts=%d preprocessing=%d",
|
|
remainingAdmissions,
|
|
remainingLeases,
|
|
remainingAttempts,
|
|
remainingPreprocessing,
|
|
)
|
|
}
|
|
|
|
postSubmission := createTask(false)
|
|
claimed, err = first.ClaimTaskPreparation(ctx, postSubmission.ID, uuid.NewString(), time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("claim post-submission task: %v", err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
INSERT INTO gateway_task_attempts (
|
|
task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
|
|
upstream_submission_status
|
|
)
|
|
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'submitting')`,
|
|
postSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
|
|
); err != nil {
|
|
t.Fatalf("create submitted attempt: %v", err)
|
|
}
|
|
_, changed, err = first.CancelTaskBeforeUpstreamSubmission(
|
|
ctx,
|
|
postSubmission.ID,
|
|
claimed.ExecutionToken,
|
|
"client disconnected after upstream submission",
|
|
)
|
|
if err != nil || changed {
|
|
t.Fatalf("post-submission cancellation changed=%v err=%v, want unchanged", changed, err)
|
|
}
|
|
|
|
queuedAtomicTask := createTask(true)
|
|
queuedHookFailure := errors.New("synthetic queued hook failure")
|
|
_, err = first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(pgx.Tx) error {
|
|
return queuedHookFailure
|
|
})
|
|
if !errors.Is(err, queuedHookFailure) {
|
|
t.Fatalf("queued hook failure = %v, want synthetic failure", err)
|
|
}
|
|
var queuedAdmissions, queuedLeases int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
|
|
(SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
|
|
queuedAtomicTask.ID,
|
|
).Scan(&queuedAdmissions, &queuedLeases); err != nil {
|
|
t.Fatalf("read rolled back queued hook state: %v", err)
|
|
}
|
|
if queuedAdmissions != 0 || queuedLeases != 0 {
|
|
t.Fatalf("failed queued hook left admissions=%d leases=%d", queuedAdmissions, queuedLeases)
|
|
}
|
|
const queuedSyntheticRiverJobID int64 = 987654320
|
|
queuedAdmission, err := first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(tx pgx.Tx) error {
|
|
_, hookErr := tx.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET river_job_id = $2
|
|
WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID)
|
|
return hookErr
|
|
})
|
|
if err != nil || queuedAdmission.Status != "waiting" {
|
|
t.Fatalf("atomic queued admission=%+v err=%v, want waiting", queuedAdmission, err)
|
|
}
|
|
var riverJobID int64
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
river_job_id,
|
|
(SELECT COUNT(*)
|
|
FROM gateway_concurrency_leases
|
|
WHERE task_id = $1::uuid AND released_at IS NULL)
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, queuedAtomicTask.ID).Scan(&riverJobID, &queuedLeases); err != nil {
|
|
t.Fatalf("read atomic queued River marker: %v", err)
|
|
}
|
|
if riverJobID != queuedSyntheticRiverJobID || queuedLeases != 0 {
|
|
t.Fatalf("queued River marker=%d leases=%d, want %d/0", riverJobID, queuedLeases, queuedSyntheticRiverJobID)
|
|
}
|
|
result, err = second.TryTaskAdmission(ctx, inputFor(queuedAtomicTask, 100, ""))
|
|
if err != nil || !result.Admitted || len(result.Leases) != 1 {
|
|
t.Fatalf("worker-time queued admission result=%+v err=%v", result, err)
|
|
}
|
|
if err := first.DeleteTaskAdmission(ctx, queuedAtomicTask.ID); err != nil {
|
|
t.Fatalf("release queued atomic task: %v", err)
|
|
}
|
|
|
|
atomicTask := createTask(true)
|
|
hookFailure := errors.New("synthetic admitted hook failure")
|
|
_, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error {
|
|
return hookFailure
|
|
})
|
|
if !errors.Is(err, hookFailure) {
|
|
t.Fatalf("admitted hook failure = %v, want synthetic failure", err)
|
|
}
|
|
var atomicAdmissions, atomicLeases int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
|
|
(SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
|
|
atomicTask.ID,
|
|
).Scan(&atomicAdmissions, &atomicLeases); err != nil {
|
|
t.Fatalf("read rolled back admitted hook state: %v", err)
|
|
}
|
|
if atomicAdmissions != 0 || atomicLeases != 0 {
|
|
t.Fatalf("failed admitted hook left admissions=%d leases=%d", atomicAdmissions, atomicLeases)
|
|
}
|
|
const syntheticRiverJobID int64 = 987654321
|
|
result, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(tx pgx.Tx) error {
|
|
_, hookErr := tx.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET river_job_id = $2
|
|
WHERE id = $1::uuid`, atomicTask.ID, syntheticRiverJobID)
|
|
return hookErr
|
|
})
|
|
if err != nil || !result.Admitted {
|
|
t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
|
|
}
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT river_job_id
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, atomicTask.ID).Scan(&riverJobID); err != nil {
|
|
t.Fatalf("read atomic River job marker: %v", err)
|
|
}
|
|
if riverJobID != syntheticRiverJobID {
|
|
t.Fatalf("atomic River job marker=%d, want %d", riverJobID, syntheticRiverJobID)
|
|
}
|
|
if err := first.DeleteTaskAdmission(ctx, atomicTask.ID); err != nil {
|
|
t.Fatalf("release atomic admitted hook task: %v", err)
|
|
}
|
|
|
|
recoveryTask := createTask(true)
|
|
result, err = first.TryTaskAdmission(ctx, inputFor(recoveryTask, 100, ""))
|
|
if err != nil || !result.Admitted {
|
|
t.Fatalf("admit interrupted recovery task: result=%+v err=%v", result, err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET status = 'running',
|
|
execution_token = $2::uuid,
|
|
execution_lease_expires_at = now() - interval '1 second'
|
|
WHERE id = $1::uuid`, recoveryTask.ID, uuid.NewString()); err != nil {
|
|
t.Fatalf("expire interrupted recovery task lease: %v", err)
|
|
}
|
|
recovery, err := first.RecoverInterruptedRuntimeState(ctx)
|
|
if err != nil {
|
|
t.Fatalf("recover interrupted admitted task: %v", err)
|
|
}
|
|
if recovery.RequeuedAsyncTasks < 1 || recovery.CleanedTaskAdmissions < 1 {
|
|
t.Fatalf("interrupted task recovery = %+v, want requeue and admission cleanup", recovery)
|
|
}
|
|
var recoveredStatus string
|
|
var recoveredExecutionToken *string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT status, execution_token::text
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, recoveryTask.ID).Scan(&recoveredStatus, &recoveredExecutionToken); err != nil {
|
|
t.Fatalf("read recovered task state: %v", err)
|
|
}
|
|
if recoveredStatus != "queued" || recoveredExecutionToken != nil {
|
|
t.Fatalf("recovered task status=%s execution_token=%v, want queued without token", recoveredStatus, recoveredExecutionToken)
|
|
}
|
|
var recoveredAdmissions, recoveredLeases int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
|
|
(SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
|
|
recoveryTask.ID,
|
|
).Scan(&recoveredAdmissions, &recoveredLeases); err != nil {
|
|
t.Fatalf("read recovered task cleanup: %v", err)
|
|
}
|
|
if recoveredAdmissions != 0 || recoveredLeases != 0 {
|
|
t.Fatalf("recovered task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
|
|
}
|
|
|
|
lockedRecoveryTask := createTask(true)
|
|
unlockedRecoveryTask := createTask(true)
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET status = 'running',
|
|
execution_token = gen_random_uuid(),
|
|
execution_lease_expires_at = now() - interval '1 second'
|
|
WHERE id = ANY($1::uuid[])`, []string{lockedRecoveryTask.ID, unlockedRecoveryTask.ID}); err != nil {
|
|
t.Fatalf("prepare locked recovery tasks: %v", err)
|
|
}
|
|
lockTx, err := second.pool.Begin(ctx)
|
|
if err != nil {
|
|
t.Fatalf("begin task row lock: %v", err)
|
|
}
|
|
var lockedTaskID string
|
|
if err := lockTx.QueryRow(ctx, `
|
|
SELECT id::text
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid
|
|
FOR UPDATE`, lockedRecoveryTask.ID).Scan(&lockedTaskID); err != nil {
|
|
rollbackTransaction(lockTx)
|
|
t.Fatalf("lock interrupted task row: %v", err)
|
|
}
|
|
recoveryCtx, recoveryCancel := context.WithTimeout(ctx, 3*time.Second)
|
|
recovery, err = first.RecoverInterruptedRuntimeState(recoveryCtx)
|
|
recoveryCancel()
|
|
if err != nil {
|
|
rollbackTransaction(lockTx)
|
|
t.Fatalf("recover unlocked task while another row is locked: %v", err)
|
|
}
|
|
if recovery.RequeuedAsyncTasks < 1 {
|
|
rollbackTransaction(lockTx)
|
|
t.Fatalf("SKIP LOCKED recovery = %+v, want at least one requeued task", recovery)
|
|
}
|
|
var lockedStatus, unlockedStatus string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT status FROM gateway_tasks WHERE id = $1::uuid),
|
|
(SELECT status FROM gateway_tasks WHERE id = $2::uuid)`,
|
|
lockedRecoveryTask.ID,
|
|
unlockedRecoveryTask.ID,
|
|
).Scan(&lockedStatus, &unlockedStatus); err != nil {
|
|
rollbackTransaction(lockTx)
|
|
t.Fatalf("read SKIP LOCKED recovery states: %v", err)
|
|
}
|
|
if lockedStatus != "running" || unlockedStatus != "queued" {
|
|
rollbackTransaction(lockTx)
|
|
t.Fatalf("SKIP LOCKED recovery statuses = %s/%s, want running/queued", lockedStatus, unlockedStatus)
|
|
}
|
|
rollbackTransaction(lockTx)
|
|
recovery, err = first.RecoverInterruptedRuntimeState(ctx)
|
|
if err != nil {
|
|
t.Fatalf("recover previously locked task: %v", err)
|
|
}
|
|
if recovery.RequeuedAsyncTasks < 1 {
|
|
t.Fatalf("unlocked recovery = %+v, want requeued task", recovery)
|
|
}
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT status
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, lockedRecoveryTask.ID).Scan(&lockedStatus); err != nil {
|
|
t.Fatalf("read unlocked recovery status: %v", err)
|
|
}
|
|
if lockedStatus != "queued" {
|
|
t.Fatalf("unlocked recovery status = %s, want queued", lockedStatus)
|
|
}
|
|
|
|
preSubmissionTask := createTask(true)
|
|
preSubmissionToken := uuid.NewString()
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET status = 'running',
|
|
execution_token = $2::uuid,
|
|
execution_lease_expires_at = now() + interval '5 minutes'
|
|
WHERE id = $1::uuid`,
|
|
preSubmissionTask.ID,
|
|
preSubmissionToken,
|
|
); err != nil {
|
|
t.Fatalf("prepare pre-submission lock timeout task: %v", err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
INSERT INTO gateway_task_attempts (
|
|
task_id, attempt_no, queue_key, status, upstream_submission_status
|
|
)
|
|
VALUES ($1::uuid, 1, 'integration-test', 'running', 'not_submitted')`,
|
|
preSubmissionTask.ID,
|
|
); err != nil {
|
|
t.Fatalf("create pre-submission attempt: %v", err)
|
|
}
|
|
requeuedTask, changed, err := first.RequeueTaskBeforeUpstreamSubmission(
|
|
ctx,
|
|
preSubmissionTask.ID,
|
|
preSubmissionToken,
|
|
time.Second,
|
|
)
|
|
if err != nil || !changed || requeuedTask.Status != "queued" {
|
|
t.Fatalf("pre-submission lock timeout requeue task=%+v changed=%v err=%v", requeuedTask, changed, err)
|
|
}
|
|
var preSubmissionAttempts int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT count(*)
|
|
FROM gateway_task_attempts
|
|
WHERE task_id = $1::uuid`, preSubmissionTask.ID).Scan(&preSubmissionAttempts); err != nil {
|
|
t.Fatalf("count cleaned pre-submission attempts: %v", err)
|
|
}
|
|
if preSubmissionAttempts != 0 {
|
|
t.Fatalf("pre-submission lock timeout left %d attempts, want 0", preSubmissionAttempts)
|
|
}
|
|
|
|
ambiguousTask := createTask(true)
|
|
ambiguousToken := uuid.NewString()
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET status = 'running',
|
|
execution_token = $2::uuid,
|
|
execution_lease_expires_at = now() + interval '5 minutes'
|
|
WHERE id = $1::uuid`,
|
|
ambiguousTask.ID,
|
|
ambiguousToken,
|
|
); err != nil {
|
|
t.Fatalf("prepare ambiguous submission task: %v", err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
INSERT INTO gateway_task_attempts (
|
|
task_id, attempt_no, queue_key, status, upstream_submission_status
|
|
)
|
|
VALUES ($1::uuid, 1, 'integration-test', 'running', 'submitting')`,
|
|
ambiguousTask.ID,
|
|
); err != nil {
|
|
t.Fatalf("create ambiguous submission attempt: %v", err)
|
|
}
|
|
_, changed, err = first.RequeueTaskBeforeUpstreamSubmission(
|
|
ctx,
|
|
ambiguousTask.ID,
|
|
ambiguousToken,
|
|
time.Second,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("guard ambiguous submission task requeue: %v", err)
|
|
}
|
|
if changed {
|
|
t.Fatal("ambiguous submission task was requeued")
|
|
}
|
|
var ambiguousStatus string
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT status
|
|
FROM gateway_tasks
|
|
WHERE id = $1::uuid`, ambiguousTask.ID).Scan(&ambiguousStatus); err != nil {
|
|
t.Fatalf("read ambiguous submission task: %v", err)
|
|
}
|
|
if ambiguousStatus != "running" {
|
|
t.Fatalf("ambiguous submission task status = %s, want running", ambiguousStatus)
|
|
}
|
|
|
|
terminalResidue := createTask(true)
|
|
result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, ""))
|
|
if err != nil || !result.Admitted {
|
|
t.Fatalf("admit terminal residue task: result=%+v err=%v", result, err)
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET status = 'failed',
|
|
error_code = 'integration_test',
|
|
finished_at = now(),
|
|
updated_at = now()
|
|
WHERE id = $1::uuid`, terminalResidue.ID); err != nil {
|
|
t.Fatalf("mark terminal residue task failed: %v", err)
|
|
}
|
|
recovery, err = first.RecoverInterruptedRuntimeState(ctx)
|
|
if err != nil {
|
|
t.Fatalf("clean terminal admission residue: %v", err)
|
|
}
|
|
if recovery.CleanedTaskAdmissions < 1 {
|
|
t.Fatalf("terminal recovery = %+v, want admission cleanup", recovery)
|
|
}
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
|
|
(SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
|
|
terminalResidue.ID,
|
|
).Scan(&recoveredAdmissions, &recoveredLeases); err != nil {
|
|
t.Fatalf("count terminal admission residue: %v", err)
|
|
}
|
|
if recoveredAdmissions != 0 || recoveredLeases != 0 {
|
|
t.Fatalf("terminal task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
|
|
}
|
|
}
|
|
|
|
func TestWorkerCapacityAllocationAndFailover(t *testing.T) {
|
|
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
|
if databaseURL == "" {
|
|
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run worker registry PostgreSQL integration tests")
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
applyOIDCJITTestMigrations(t, ctx, databaseURL)
|
|
db, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
defer db.Close()
|
|
var databaseName string
|
|
if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
|
|
t.Fatalf("read test database name: %v", err)
|
|
}
|
|
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
|
t.Fatalf("refusing to use non-test database %q", databaseName)
|
|
}
|
|
|
|
prefix := "queue-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
|
|
firstID := prefix + "-a"
|
|
secondID := prefix + "-b"
|
|
defer db.pool.Exec(context.Background(), `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{firstID, secondID})
|
|
|
|
first, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
|
|
if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
|
|
t.Fatalf("first allocation = %+v, err=%v", first, err)
|
|
}
|
|
second, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: secondID, DesiredCapacity: 5})
|
|
if err != nil || second.Allocated != 2 || second.ActiveInstances != 2 {
|
|
t.Fatalf("second allocation = %+v, err=%v", second, err)
|
|
}
|
|
first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
|
|
if err != nil || first.Allocated != 3 || first.ActiveInstances != 2 {
|
|
t.Fatalf("rebalanced first allocation = %+v, err=%v", first, err)
|
|
}
|
|
if _, err := db.pool.Exec(ctx, `
|
|
UPDATE gateway_worker_instances
|
|
SET heartbeat_at = now() - $2::interval
|
|
WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).String()); err != nil {
|
|
t.Fatalf("expire second heartbeat: %v", err)
|
|
}
|
|
first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
|
|
if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
|
|
t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
|
|
}
|
|
|
|
first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
|
|
InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 3,
|
|
})
|
|
if err != nil || first.Allocated != 3 || first.GlobalAllocated != 3 {
|
|
t.Fatalf("bounded single-worker allocation = %+v, err=%v", first, err)
|
|
}
|
|
second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
|
|
InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 2,
|
|
})
|
|
if err != nil || second.Allocated != 2 || second.GlobalAllocated != 5 || second.ActiveInstances != 2 {
|
|
t.Fatalf("bounded two-worker allocation = %+v, err=%v", second, err)
|
|
}
|
|
}
|