fix(api): 分离异步队列入队与执行开关

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 ./...。
This commit is contained in:
2026-07-29 03:03:14 +08:00
parent 0bef13fbd7
commit 9a3cc582ce
3 changed files with 106 additions and 4 deletions
+5 -2
View File
@@ -122,8 +122,11 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
if cfg.AsyncQueueWorkerEnabled { if cfg.AsyncQueueWorkerEnabled {
server.runner.StartAsyncQueueWorker(ctx) server.runner.StartAsyncQueueWorker(ctx)
} else if logger != nil { } else {
logger.Info("asynchronous queue worker disabled for this process") server.runner.StartAsyncQueueClient(ctx)
if logger != nil {
logger.Info("asynchronous queue worker disabled for this process")
}
} }
server.runner.StartBillingSettlementWorker(ctx) server.runner.StartBillingSettlementWorker(ctx)
server.runner.StartTaskHistoryWorkers(ctx) server.runner.StartTaskHistoryWorkers(ctx)
@@ -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)
}
}
+27 -2
View File
@@ -78,13 +78,20 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
} }
func (s *Service) StartAsyncQueueWorker(ctx context.Context) { 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) s.logger.Error("start river async queue failed", "error", err)
panic(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()) driver := riverpgxv5.New(s.store.Pool())
migrator, err := rivermigrate.New(driver, nil) migrator, err := rivermigrate.New(driver, nil)
if err != nil { if err != nil {
@@ -102,6 +109,23 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
if err != nil { if err != nil {
return err 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) snapshot, err := s.loadAsyncWorkerCapacity(ctx)
if err != nil { if err != nil {
return fmt.Errorf("calculate initial async worker capacity: %w", err) 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) { func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
<-ctx.Done() <-ctx.Done()
s.riverMu.Lock() s.riverMu.Lock()
s.riverControlClient = nil
clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients)) clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
if s.riverExecutionClient != nil { if s.riverExecutionClient != nil {
clients = append(clients, s.riverExecutionClient) clients = append(clients, s.riverExecutionClient)