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:
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user