wangbo
38f87b6970
队列恢复后，仍在运行标记中的 River job 可能保留 waiting admission，形成全局 FIFO 队头阻塞；同时旧执行在租约被接管后仍可能创建 attempt 或切换到 submitting。 新增 stale admission 自动让出机制，并用当前 execution token 对 attempt 创建和上游提交状态切换做 fencing。任务、River job 和结算状态均不在让出流程中改写。 验证：go vet ./...；env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1；真实 PostgreSQL 集成测试覆盖 stale admission 让出与旧 token 提交拦截。
275 lines
9.2 KiB
Go
275 lines
9.2 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
)
|
|
|
|
func TestAsyncQueueClientEnqueuesWithoutExecutionWorker(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 the queue client integration test")
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
db, err := store.Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
t.Cleanup(db.Close)
|
|
|
|
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
|
task, err := db.CreateTask(ctx, store.CreateTaskInput{
|
|
Kind: "images.edits",
|
|
Model: "queue-client-" + suffix,
|
|
Request: map[string]any{"prompt": "enqueue without execution worker"},
|
|
Async: true,
|
|
RunMode: "simulation",
|
|
}, &auth.User{ID: "queue-client-" + suffix, Source: "gateway"})
|
|
if err != nil {
|
|
t.Fatalf("create task: %v", err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.Pool().Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id = $1::uuid`, task.ID)
|
|
})
|
|
|
|
service := New(config.Config{AppEnv: "test"}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
|
service.StartAsyncQueueClient(ctx)
|
|
if err := service.EnqueueAsyncTask(ctx, task); err != nil {
|
|
t.Fatalf("enqueue task without execution worker: %v", err)
|
|
}
|
|
|
|
queued, err := db.GetTask(ctx, task.ID)
|
|
if err != nil {
|
|
t.Fatalf("load queued task: %v", err)
|
|
}
|
|
if queued.RiverJobID <= 0 {
|
|
t.Fatal("queue client did not persist a River job ID")
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.Pool().Exec(context.Background(), `DELETE FROM river_job WHERE id = $1`, queued.RiverJobID)
|
|
})
|
|
|
|
var attemptedByCount int
|
|
if err := db.Pool().QueryRow(ctx, `
|
|
SELECT COALESCE(cardinality(attempted_by), 0)
|
|
FROM river_job
|
|
WHERE id = $1`, queued.RiverJobID).Scan(&attemptedByCount); err != nil {
|
|
t.Fatalf("load River job: %v", err)
|
|
}
|
|
if attemptedByCount != 0 {
|
|
t.Fatalf("control-only queue client executed the job: attempted_by=%d", attemptedByCount)
|
|
}
|
|
}
|
|
|
|
func TestRecoverOrphanedAsyncRiverJobAndYieldStaleAdmission(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 the orphaned River job integration test")
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
db, err := store.Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
t.Cleanup(db.Close)
|
|
service := New(config.Config{AppEnv: "test"}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
|
service.StartAsyncQueueClient(ctx)
|
|
|
|
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
|
createQueuedTask := func(label string) store.GatewayTask {
|
|
t.Helper()
|
|
task, createErr := db.CreateTask(ctx, store.CreateTaskInput{
|
|
Kind: "images.edits",
|
|
Model: "orphan-recovery-" + label + "-" + suffix,
|
|
Request: map[string]any{"prompt": "recover orphaned River job"},
|
|
Async: true,
|
|
RunMode: "simulation",
|
|
}, &auth.User{ID: "orphan-recovery-" + suffix, Source: "gateway"})
|
|
if createErr != nil {
|
|
t.Fatalf("create %s task: %v", label, createErr)
|
|
}
|
|
if enqueueErr := service.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
|
|
t.Fatalf("enqueue %s task: %v", label, enqueueErr)
|
|
}
|
|
queued, getErr := db.GetTask(ctx, task.ID)
|
|
if getErr != nil {
|
|
t.Fatalf("load %s task: %v", label, getErr)
|
|
}
|
|
return queued
|
|
}
|
|
|
|
orphaned := createQueuedTask("orphaned")
|
|
protected := createQueuedTask("protected")
|
|
yielded := createQueuedTask("yielded")
|
|
staleWorkerID := "orphan-recovery-stale-" + suffix
|
|
activeWorkerID := "orphan-recovery-active-" + suffix
|
|
t.Cleanup(func() {
|
|
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cleanupCancel()
|
|
_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{staleWorkerID, activeWorkerID})
|
|
_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM river_job WHERE id = ANY($1::bigint[])`, []int64{orphaned.RiverJobID, protected.RiverJobID, yielded.RiverJobID})
|
|
_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, []string{orphaned.ID, protected.ID, yielded.ID})
|
|
})
|
|
|
|
if _, err := db.Pool().Exec(ctx, `
|
|
INSERT INTO gateway_worker_instances (
|
|
instance_id, status, desired_capacity, allocated_capacity, heartbeat_at, updated_at
|
|
)
|
|
VALUES
|
|
($1, 'active', 1, 0, now() - interval '2 minutes', now() - interval '2 minutes'),
|
|
($2, 'active', 1, 1, now(), now())`,
|
|
staleWorkerID,
|
|
activeWorkerID,
|
|
); err != nil {
|
|
t.Fatalf("seed worker heartbeats: %v", err)
|
|
}
|
|
if _, err := db.Pool().Exec(ctx, `
|
|
UPDATE river_job
|
|
SET state = 'running',
|
|
attempt = 1,
|
|
attempted_at = now() - interval '1 minute',
|
|
attempted_by = ARRAY[
|
|
CASE id
|
|
WHEN $1::bigint THEN $3
|
|
WHEN $2::bigint THEN $4
|
|
WHEN $5::bigint THEN $4
|
|
END
|
|
]::text[]
|
|
WHERE id = ANY($6::bigint[])`,
|
|
orphaned.RiverJobID,
|
|
protected.RiverJobID,
|
|
staleWorkerID+"-exec-1-test",
|
|
activeWorkerID+"-exec-1-test",
|
|
yielded.RiverJobID,
|
|
[]int64{orphaned.RiverJobID, protected.RiverJobID, yielded.RiverJobID},
|
|
); err != nil {
|
|
t.Fatalf("mark River jobs running: %v", err)
|
|
}
|
|
var platformID, platformModelID string
|
|
if err := db.Pool().QueryRow(ctx, `
|
|
SELECT platform.id::text, model.id::text
|
|
FROM integration_platforms platform
|
|
JOIN platform_models model ON model.platform_id = platform.id
|
|
ORDER BY platform.priority ASC, model.created_at ASC
|
|
LIMIT 1`).Scan(&platformID, &platformModelID); err != nil {
|
|
t.Fatalf("load admission binding: %v", err)
|
|
}
|
|
if _, err := db.Pool().Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET execution_lease_expires_at = now() - interval '1 minute'
|
|
WHERE id = $1::uuid`,
|
|
yielded.ID,
|
|
); err != nil {
|
|
t.Fatalf("expire yielded task execution lease: %v", err)
|
|
}
|
|
if _, err := db.Pool().Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET execution_lease_expires_at = now() + interval '5 minutes'
|
|
WHERE id = $1::uuid`,
|
|
protected.ID,
|
|
); err != nil {
|
|
t.Fatalf("renew protected task execution lease: %v", err)
|
|
}
|
|
if _, err := db.Pool().Exec(ctx, `
|
|
INSERT INTO gateway_task_admissions (
|
|
task_id, platform_id, platform_model_id, queue_key, mode, status,
|
|
priority, enqueued_at, wait_deadline_at
|
|
)
|
|
VALUES (
|
|
$1::uuid, $2::uuid, $3::uuid, 'integration-test', 'async', 'waiting',
|
|
100, now() - interval '1 minute', now() + interval '10 minutes'
|
|
)`,
|
|
yielded.ID,
|
|
platformID,
|
|
platformModelID,
|
|
); err != nil {
|
|
t.Fatalf("seed stale async admission: %v", err)
|
|
}
|
|
if _, err := db.Pool().Exec(ctx, `
|
|
INSERT INTO gateway_task_admissions (
|
|
task_id, platform_id, platform_model_id, queue_key, mode, status,
|
|
priority, enqueued_at, wait_deadline_at
|
|
)
|
|
VALUES (
|
|
$1::uuid, $2::uuid, $3::uuid, 'integration-test', 'async', 'waiting',
|
|
100, now(), now() + interval '10 minutes'
|
|
)`,
|
|
protected.ID,
|
|
platformID,
|
|
platformModelID,
|
|
); err != nil {
|
|
t.Fatalf("seed protected async admission: %v", err)
|
|
}
|
|
yieldedCount, err := db.YieldStaleAsyncTaskAdmissions(ctx, 30*time.Second, 10)
|
|
if err != nil {
|
|
t.Fatalf("yield stale async task admission: %v", err)
|
|
}
|
|
if yieldedCount != 1 {
|
|
t.Fatalf("yielded admissions=%d, want 1", yieldedCount)
|
|
}
|
|
var yieldedAdmissions, protectedAdmissions int
|
|
if err := db.Pool().QueryRow(ctx, `
|
|
SELECT
|
|
count(*) FILTER (WHERE task_id = $1::uuid),
|
|
count(*) FILTER (WHERE task_id = $2::uuid)
|
|
FROM gateway_task_admissions
|
|
WHERE task_id = ANY($3::uuid[])`,
|
|
yielded.ID,
|
|
protected.ID,
|
|
[]string{yielded.ID, protected.ID},
|
|
).Scan(&yieldedAdmissions, &protectedAdmissions); err != nil {
|
|
t.Fatalf("count yielded admissions: %v", err)
|
|
}
|
|
var yieldedState string
|
|
if err := db.Pool().QueryRow(ctx, `
|
|
SELECT state::text
|
|
FROM river_job
|
|
WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedState); err != nil {
|
|
t.Fatalf("read yielded River job state: %v", err)
|
|
}
|
|
if yieldedAdmissions != 0 || protectedAdmissions != 1 || yieldedState != "running" {
|
|
t.Fatalf(
|
|
"yielded/protected admissions=%d/%d River state=%s, want 0/1/running",
|
|
yieldedAdmissions,
|
|
protectedAdmissions,
|
|
yieldedState,
|
|
)
|
|
}
|
|
|
|
recovered, err := db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
|
|
if err != nil {
|
|
t.Fatalf("recover orphaned River jobs: %v", err)
|
|
}
|
|
if recovered != 1 {
|
|
t.Fatalf("recovered jobs=%d, want 1", recovered)
|
|
}
|
|
var orphanedState, protectedState string
|
|
if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, orphaned.RiverJobID).Scan(&orphanedState); err != nil {
|
|
t.Fatalf("read orphaned River job state: %v", err)
|
|
}
|
|
if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, protected.RiverJobID).Scan(&protectedState); err != nil {
|
|
t.Fatalf("read protected River job state: %v", err)
|
|
}
|
|
if orphanedState != "retryable" || protectedState != "running" {
|
|
t.Fatalf("River states orphaned=%s protected=%s, want retryable/running", orphanedState, protectedState)
|
|
}
|
|
recovered, err = db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
|
|
if err != nil || recovered != 0 {
|
|
t.Fatalf("second orphan recovery count=%d err=%v, want idempotent zero", recovered, err)
|
|
}
|
|
}
|