From 9a3cc582ce8a28d3252e8a684a891b7812191555 Mon Sep 17 00:00:00 2001 From: wangbo Date: Wed, 29 Jul 2026 03:03:14 +0800 Subject: [PATCH] =?UTF-8?q?fix(api):=20=E5=88=86=E7=A6=BB=E5=BC=82?= =?UTF-8?q?=E6=AD=A5=E9=98=9F=E5=88=97=E5=85=A5=E9=98=9F=E4=B8=8E=E6=89=A7?= =?UTF-8?q?=E8=A1=8C=E5=BC=80=E5=85=B3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED=false 时仍初始化 River schema 与控制客户端，允许 HTTP 请求持久化并入队；仅跳过执行客户端和动态容量 Worker。关闭上下文时同步清理控制客户端引用。\n\n新增数据库集成测试覆盖无执行 Worker 时成功创建 River job 且 attempted_by 为空。验证：gofmt、go test ./... -count=1、go vet ./...。 --- apps/api/internal/httpapi/server.go | 7 +- apps/api/internal/runner/queue_client_test.go | 74 +++++++++++++++++++ apps/api/internal/runner/queue_worker.go | 29 +++++++- 3 files changed, 106 insertions(+), 4 deletions(-) create mode 100644 apps/api/internal/runner/queue_client_test.go diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index e0682ee..d5ed588 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -122,8 +122,11 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey if cfg.AsyncQueueWorkerEnabled { server.runner.StartAsyncQueueWorker(ctx) - } else if logger != nil { - logger.Info("asynchronous queue worker disabled for this process") + } else { + server.runner.StartAsyncQueueClient(ctx) + if logger != nil { + logger.Info("asynchronous queue worker disabled for this process") + } } server.runner.StartBillingSettlementWorker(ctx) server.runner.StartTaskHistoryWorkers(ctx) diff --git a/apps/api/internal/runner/queue_client_test.go b/apps/api/internal/runner/queue_client_test.go new file mode 100644 index 0000000..9820f79 --- /dev/null +++ b/apps/api/internal/runner/queue_client_test.go @@ -0,0 +1,74 @@ +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 cardinality(attempted_by) +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) + } +} diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index 72e0b9b..3de7bfa 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -78,13 +78,20 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs } func (s *Service) StartAsyncQueueWorker(ctx context.Context) { - if err := s.startRiverQueue(ctx); err != nil { + if err := s.startRiverQueue(ctx, true); err != nil { s.logger.Error("start river async queue failed", "error", err) panic(err) } } -func (s *Service) startRiverQueue(ctx context.Context) error { +func (s *Service) StartAsyncQueueClient(ctx context.Context) { + if err := s.startRiverQueue(ctx, false); err != nil { + s.logger.Error("start river async queue client failed", "error", err) + panic(err) + } +} + +func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error { driver := riverpgxv5.New(s.store.Pool()) migrator, err := rivermigrate.New(driver, nil) if err != nil { @@ -102,6 +109,23 @@ func (s *Service) startRiverQueue(ctx context.Context) error { if err != nil { return err } + if !workerEnabled { + s.riverMu.Lock() + s.riverControlClient = controlClient + s.riverExecutionClient = nil + s.riverWorkerCapacity = 0 + s.riverDrainingClients = make(map[asyncExecutionClient]struct{}) + s.riverMu.Unlock() + if err := s.recoverAsyncRiverJobs(ctx); err != nil { + s.riverMu.Lock() + s.riverControlClient = nil + s.riverMu.Unlock() + return err + } + go s.stopAsyncWorkersOnShutdown(ctx) + s.logger.Info("river async queue client initialized without execution workers") + return nil + } snapshot, err := s.loadAsyncWorkerCapacity(ctx) if err != nil { return fmt.Errorf("calculate initial async worker capacity: %w", err) @@ -271,6 +295,7 @@ func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) { func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) { <-ctx.Done() s.riverMu.Lock() + s.riverControlClient = nil clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients)) if s.riverExecutionClient != nil { clients = append(clients, s.riverExecutionClient)