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