wangbo
5dd7765ae3
跨地域同步复制下逐任务准入会为每个 River Job 支付一次事务提交 RTT，导致 P24 实际只能维持少量运行任务。调度器现在按全局容量窗口准备 FIFO 批次，在同一 PostgreSQL 事务内逐项校验队首、创建租约并插入唯一 River Job，一次提交即可填满可用槽；任一 Hook 失败时整批原子回滚。\n\n验收压力采样同时排除正在终止的 Ready Pod，避免滚动切换期间误连已移除容器。\n\n验证：真实 PostgreSQL 批量提交及整批回滚集成测试、Go 全量测试、go vet、runner/store race、gofmt、bash -n、ShellCheck、迁移安全检查通过。
1129 lines
41 KiB
Go
1129 lines
41 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"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)
|
|
}
|
|
waitingAsyncTaskIDs, err := first.ListWaitingAsyncAdmissionTaskIDs(ctx, 1000)
|
|
if err != nil {
|
|
t.Fatalf("list waiting async tasks with missing River job: %v", err)
|
|
}
|
|
waitingAsyncFound := false
|
|
for _, taskID := range waitingAsyncTaskIDs {
|
|
if taskID == queuedAtomicTask.ID {
|
|
waitingAsyncFound = true
|
|
break
|
|
}
|
|
}
|
|
if !waitingAsyncFound {
|
|
t.Fatal("waiting async task with missing River job was not dispatchable")
|
|
}
|
|
recoverableTasks, err := first.ListRecoverableAsyncTasks(ctx, 1000)
|
|
if err != nil {
|
|
t.Fatalf("list recoverable tasks while admission is waiting: %v", err)
|
|
}
|
|
for _, item := range recoverableTasks {
|
|
if item.ID == queuedAtomicTask.ID {
|
|
t.Fatal("generic River recovery claimed a task owned by the async admission dispatcher")
|
|
}
|
|
}
|
|
recoveryGraceTask := createTask(true)
|
|
recoverableTasks, err = first.ListRecoverableAsyncTasks(ctx, 1000)
|
|
if err != nil {
|
|
t.Fatalf("list fresh recoverable tasks: %v", err)
|
|
}
|
|
for _, item := range recoverableTasks {
|
|
if item.ID == recoveryGraceTask.ID {
|
|
t.Fatal("generic River recovery claimed a task still preparing admission")
|
|
}
|
|
}
|
|
if _, err := first.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET created_at = now() - interval '6 minutes'
|
|
WHERE id = $1::uuid`, recoveryGraceTask.ID); err != nil {
|
|
t.Fatalf("age task beyond River recovery grace period: %v", err)
|
|
}
|
|
recoverableTasks, err = first.ListRecoverableAsyncTasks(ctx, 1000)
|
|
if err != nil {
|
|
t.Fatalf("list aged recoverable tasks: %v", err)
|
|
}
|
|
recoveryGraceFound := false
|
|
for _, item := range recoverableTasks {
|
|
if item.ID == recoveryGraceTask.ID {
|
|
recoveryGraceFound = true
|
|
break
|
|
}
|
|
}
|
|
if !recoveryGraceFound {
|
|
t.Fatal("generic River recovery did not claim a task beyond the recovery grace period")
|
|
}
|
|
|
|
batchScope := AdmissionScope{
|
|
ScopeType: "worker_capacity",
|
|
ScopeKey: "batch-" + suffix,
|
|
ScopeName: "batch capacity",
|
|
ConcurrentLimit: 3,
|
|
Amount: 1,
|
|
LeaseTTLSeconds: 120,
|
|
QueueLimit: 100,
|
|
MaxWaitSeconds: 600,
|
|
}
|
|
batchTasks := []GatewayTask{createTask(true), createTask(true), createTask(true)}
|
|
batchInputs := make([]TaskAdmissionInput, 0, len(batchTasks))
|
|
batchJobIDs := make(map[string]int64, len(batchTasks))
|
|
for index, task := range batchTasks {
|
|
input := inputFor(task, 100+index, "")
|
|
input.Scopes = []AdmissionScope{batchScope}
|
|
if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
|
|
t.Fatalf("queue batch task %d: %v", index, err)
|
|
}
|
|
batchInputs = append(batchInputs, input)
|
|
batchJobIDs[task.ID] = int64(987654400 + index)
|
|
}
|
|
batchOutcomes, err := first.TryTaskAdmissionBatchWithAdmittedHook(
|
|
ctx,
|
|
batchInputs,
|
|
func(tx pgx.Tx, input TaskAdmissionInput) error {
|
|
_, hookErr := tx.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET river_job_id = $2
|
|
WHERE id = $1::uuid`, input.TaskID, batchJobIDs[input.TaskID])
|
|
return hookErr
|
|
},
|
|
)
|
|
if err != nil || len(batchOutcomes) != len(batchTasks) {
|
|
t.Fatalf("batch admission outcomes=%+v err=%v", batchOutcomes, err)
|
|
}
|
|
for index, outcome := range batchOutcomes {
|
|
if outcome.Err != nil || !outcome.Result.Admitted || !outcome.Result.NewlyAdmitted ||
|
|
len(outcome.Result.Leases) != 1 {
|
|
t.Fatalf("batch outcome %d=%+v", index, outcome)
|
|
}
|
|
}
|
|
var batchAdmissions, batchLeases, batchRiverJobs int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT count(*)
|
|
FROM gateway_task_admissions
|
|
WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'),
|
|
(SELECT count(*)
|
|
FROM gateway_concurrency_leases
|
|
WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL AND expires_at > now()),
|
|
(SELECT count(*)
|
|
FROM gateway_tasks
|
|
WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
|
|
[]string{batchTasks[0].ID, batchTasks[1].ID, batchTasks[2].ID},
|
|
).Scan(&batchAdmissions, &batchLeases, &batchRiverJobs); err != nil {
|
|
t.Fatalf("read committed batch admission: %v", err)
|
|
}
|
|
if batchAdmissions != 3 || batchLeases != 3 || batchRiverJobs != 3 {
|
|
t.Fatalf(
|
|
"batch admissions=%d leases=%d River jobs=%d, want 3/3/3",
|
|
batchAdmissions,
|
|
batchLeases,
|
|
batchRiverJobs,
|
|
)
|
|
}
|
|
for _, task := range batchTasks {
|
|
if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
|
|
t.Fatalf("release batch task %s: %v", task.ID, err)
|
|
}
|
|
}
|
|
|
|
rollbackScope := batchScope
|
|
rollbackScope.ScopeKey = "batch-rollback-" + suffix
|
|
rollbackScope.ConcurrentLimit = 2
|
|
rollbackTasks := []GatewayTask{createTask(true), createTask(true)}
|
|
rollbackInputs := make([]TaskAdmissionInput, 0, len(rollbackTasks))
|
|
for index, task := range rollbackTasks {
|
|
input := inputFor(task, 200+index, "")
|
|
input.Scopes = []AdmissionScope{rollbackScope}
|
|
if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
|
|
t.Fatalf("queue rollback batch task %d: %v", index, err)
|
|
}
|
|
rollbackInputs = append(rollbackInputs, input)
|
|
}
|
|
batchHookFailure := errors.New("synthetic batch hook failure")
|
|
_, err = first.TryTaskAdmissionBatchWithAdmittedHook(
|
|
ctx,
|
|
rollbackInputs,
|
|
func(tx pgx.Tx, input TaskAdmissionInput) error {
|
|
if input.TaskID == rollbackTasks[1].ID {
|
|
return batchHookFailure
|
|
}
|
|
_, hookErr := tx.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET river_job_id = 987654499
|
|
WHERE id = $1::uuid`, input.TaskID)
|
|
return hookErr
|
|
},
|
|
)
|
|
if !errors.Is(err, batchHookFailure) {
|
|
t.Fatalf("batch hook failure=%v, want synthetic failure", err)
|
|
}
|
|
var rollbackWaiting, rollbackLeases, rollbackRiverJobs int
|
|
if err := first.pool.QueryRow(ctx, `
|
|
SELECT
|
|
(SELECT count(*)
|
|
FROM gateway_task_admissions
|
|
WHERE task_id = ANY($1::uuid[]) AND status = 'waiting'),
|
|
(SELECT count(*)
|
|
FROM gateway_concurrency_leases
|
|
WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL),
|
|
(SELECT count(*)
|
|
FROM gateway_tasks
|
|
WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
|
|
[]string{rollbackTasks[0].ID, rollbackTasks[1].ID},
|
|
).Scan(&rollbackWaiting, &rollbackLeases, &rollbackRiverJobs); err != nil {
|
|
t.Fatalf("read rolled back batch admission: %v", err)
|
|
}
|
|
if rollbackWaiting != 2 || rollbackLeases != 0 || rollbackRiverJobs != 0 {
|
|
t.Fatalf(
|
|
"rolled back batch waiting=%d leases=%d River jobs=%d, want 2/0/0",
|
|
rollbackWaiting,
|
|
rollbackLeases,
|
|
rollbackRiverJobs,
|
|
)
|
|
}
|
|
for _, task := range rollbackTasks {
|
|
if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
|
|
t.Fatalf("release rollback batch task %s: %v", task.ID, err)
|
|
}
|
|
}
|
|
|
|
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 TestWaitingAdmissionLeadersDoNotCrossBlockIndependentCandidates(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(), 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)
|
|
}
|
|
|
|
suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
|
|
platform, err := db.CreatePlatform(ctx, CreatePlatformInput{
|
|
Provider: "leader-test",
|
|
PlatformKey: "leader-test-" + suffix,
|
|
Name: "Admission Leader Test " + suffix,
|
|
AuthType: "none",
|
|
Status: "enabled",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("create platform: %v", err)
|
|
}
|
|
modelIDs := make([]string, 0, 2)
|
|
for index := 0; index < 2; index++ {
|
|
modelName := fmt.Sprintf("leader-model-%d-%s", index, suffix)
|
|
var modelID string
|
|
if err := db.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, modelName).Scan(&modelID); err != nil {
|
|
t.Fatalf("create platform model %d: %v", index, err)
|
|
}
|
|
modelIDs = append(modelIDs, modelID)
|
|
}
|
|
group, err := db.CreateUserGroup(ctx, UserGroupInput{
|
|
GroupKey: "leader-test-" + suffix,
|
|
Name: "Admission Leader Test " + suffix,
|
|
Status: "active",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("create user group: %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 {
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
|
|
}
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_user_groups WHERE id = $1::uuid`, group.ID)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = ANY($1::uuid[])`, modelIDs)
|
|
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
|
|
})
|
|
createTask := func(modelIndex int) GatewayTask {
|
|
t.Helper()
|
|
task, createErr := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generate",
|
|
Model: fmt.Sprintf("leader-model-%d-%s", modelIndex, suffix),
|
|
RunMode: "production",
|
|
Request: map[string]any{"prompt": "admission leader test"},
|
|
}, &auth.User{ID: "leader-test-user-" + suffix, Source: "gateway"})
|
|
if createErr != nil {
|
|
t.Fatalf("create task: %v", createErr)
|
|
}
|
|
taskIDs = append(taskIDs, task.ID)
|
|
return task
|
|
}
|
|
tasks := []GatewayTask{createTask(0), createTask(0), createTask(1), createTask(1)}
|
|
insertWaiting := func(task GatewayTask, modelID string, groupID string, enqueuedAt time.Time) {
|
|
t.Helper()
|
|
if _, insertErr := db.pool.Exec(ctx, `
|
|
INSERT INTO gateway_task_admissions (
|
|
task_id, platform_id, platform_model_id, user_group_id, queue_key, mode,
|
|
status, priority, enqueued_at, wait_deadline_at, waiter_id,
|
|
waiter_lease_expires_at
|
|
)
|
|
VALUES (
|
|
$1::uuid, $2::uuid, $3::uuid, NULLIF($4, '')::uuid, $5, 'sync',
|
|
'waiting', 100, $6, now() + interval '10 minutes', $7,
|
|
now() + interval '15 seconds'
|
|
)`,
|
|
task.ID,
|
|
platform.ID,
|
|
modelID,
|
|
groupID,
|
|
"leader-test:"+modelID,
|
|
enqueuedAt,
|
|
"waiter-"+task.ID,
|
|
); insertErr != nil {
|
|
t.Fatalf("insert waiting admission: %v", insertErr)
|
|
}
|
|
}
|
|
baseTime := time.Now().Add(-time.Minute)
|
|
insertWaiting(tasks[0], modelIDs[0], "", baseTime)
|
|
insertWaiting(tasks[1], modelIDs[0], "", baseTime.Add(time.Second))
|
|
insertWaiting(tasks[2], modelIDs[1], "", baseTime.Add(2*time.Second))
|
|
insertWaiting(tasks[3], modelIDs[1], "", baseTime.Add(3*time.Second))
|
|
leaders, err := db.ListWaitingTaskAdmissionIDs(ctx, 256)
|
|
if err != nil {
|
|
t.Fatalf("list independent candidate leaders: %v", err)
|
|
}
|
|
leaderSet := make(map[string]bool, len(leaders))
|
|
for _, taskID := range leaders {
|
|
leaderSet[taskID] = true
|
|
}
|
|
if !leaderSet[tasks[0].ID] || !leaderSet[tasks[2].ID] {
|
|
t.Fatalf("independent candidate leaders = %v, want %s and %s", leaders, tasks[0].ID, tasks[2].ID)
|
|
}
|
|
if leaderSet[tasks[1].ID] || leaderSet[tasks[3].ID] {
|
|
t.Fatalf("non-head independent candidate task was selected: %v", leaders)
|
|
}
|
|
|
|
if _, err := db.pool.Exec(ctx, `
|
|
DELETE FROM gateway_task_admissions
|
|
WHERE task_id = ANY($1::uuid[])`, taskIDs); err != nil {
|
|
t.Fatalf("clear independent admissions: %v", err)
|
|
}
|
|
insertWaiting(tasks[0], modelIDs[0], group.ID, baseTime)
|
|
insertWaiting(tasks[2], modelIDs[1], group.ID, baseTime.Add(time.Second))
|
|
leaders, err = db.ListWaitingTaskAdmissionIDs(ctx, 256)
|
|
if err != nil {
|
|
t.Fatalf("list grouped candidate leaders: %v", err)
|
|
}
|
|
leaderSet = make(map[string]bool, len(leaders))
|
|
for _, taskID := range leaders {
|
|
leaderSet[taskID] = true
|
|
}
|
|
if !leaderSet[tasks[0].ID] || leaderSet[tasks[2].ID] {
|
|
t.Fatalf("grouped candidate leaders = %v, want only group head %s", leaders, tasks[0].ID)
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|