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
+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) {
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)