wangbo
f3dd7cd262
新增基于任务状态、执行租约和 Worker 心跳的孤儿 River Job 精确识别，每 15 秒将确认失活的 Job 恢复为可重试，避免等待 River 一小时通用救援窗口。运行时恢复同步清除中断及终态任务的 admission、执行令牌和并发租约，防止状态残留污染队列指标。正常运行中的长任务仍由 running 状态和活跃心跳保护。验证通过完整 Go 测试、go vet、竞态检查、迁移安全检查和 PostgreSQL 18 集成测试。
577 lines
22 KiB
Go
577 lines
22 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)
|
|
}
|
|
|
|
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)
|
|
}
|
|
var riverJobID int64
|
|
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)
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|