Files
easyai-ai-gateway/apps/api/internal/runner/queue_worker.go
T
wangbo fa818e9ebf fix(queue): 重新调度终态 River 任务
等待异步准入的任务只跳过仍处于活跃状态的 River Job;missing 或 terminal Job 交由准入调度器重新创建,避免非空 river_job_id 永久阻断旧任务。\n\n通用 River 恢复排除等待准入任务,消除双恢复循环与误导日志;事务回滚继续使用独立有界上下文。\n\n验证:真实 PostgreSQL 下异步 Worker 准入验收和跨 Store 队列集成测试均通过。
2026-07-30 17:42:33 +08:00

689 lines
21 KiB
Go

package runner
import (
"context"
"errors"
"fmt"
"os"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/riverqueue/river"
"github.com/riverqueue/river/riverdriver/riverpgxv5"
"github.com/riverqueue/river/rivermigrate"
"github.com/riverqueue/river/rivertype"
)
const (
asyncTaskQueueName = "gateway_tasks"
orphanedRiverJobScanInterval = 15 * time.Second
orphanedRiverJobWorkerStaleAfter = 45 * time.Second
orphanedRiverJobRecoveryBatchSize = 100
asyncQueueStartupMaxAttempts = 3
asyncQueueStartupRetryDelay = 2 * time.Second
)
type asyncTaskArgs struct {
TaskID string `json:"task_id" river:"unique"`
}
type asyncExecutionClient interface {
Start(context.Context) error
Stop(context.Context) error
StopAndCancel(context.Context) error
}
func (asyncTaskArgs) Kind() string { return "gateway_task_run" }
type asyncTaskWorker struct {
river.WorkerDefaults[asyncTaskArgs]
service *Service
}
func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs]) error {
task, err := w.service.store.GetTask(ctx, job.Args.TaskID)
if err != nil {
return err
}
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
return nil
}
executionToken := uuid.NewString()
result, runErr := w.service.executeWithToken(ctx, task, authUserFromTask(task), nil, executionToken)
if runErr == nil {
w.service.logger.Debug("river async task completed", "taskID", task.ID, "status", result.Task.Status, "riverJobID", job.ID)
return nil
}
if errors.Is(runErr, store.ErrTaskExecutionLeaseUnavailable) {
w.service.logger.Debug("river async task execution lease already held", "taskID", task.ID, "riverJobID", job.ID)
return nil
}
if errors.Is(runErr, store.ErrTaskExecutionManualReview) {
w.service.logger.Warn("river async task moved to manual review after ambiguous upstream submission", "taskID", task.ID, "riverJobID", job.ID)
return nil
}
var queuedErr *TaskQueuedError
if errors.As(runErr, &queuedErr) {
return river.JobSnooze(queuedErr.Delay)
}
if ctx.Err() != nil {
task.ExecutionToken = executionToken
queued, queueErr := w.service.requeueInterruptedAsyncTask(context.WithoutCancel(ctx), task)
if queueErr != nil {
return queueErr
}
w.service.logger.Debug("river async task interrupted and requeued", "taskID", task.ID, "status", queued.Status, "riverJobID", job.ID)
return river.JobSnooze(0)
}
if store.IsPostgresLockTimeout(runErr) {
queued, changed, queueErr := w.service.store.RequeueTaskBeforeUpstreamSubmission(
context.WithoutCancel(ctx),
task.ID,
executionToken,
time.Second,
)
if queueErr != nil {
return queueErr
}
if changed {
if eventErr := w.service.emit(
context.WithoutCancel(ctx),
task.ID,
"task.queued",
"queued",
"database_lock_timeout",
0.2,
"async task queued after database lock timeout",
map[string]any{"code": "database_lock_timeout"},
task.RunMode == "simulation",
); eventErr != nil {
w.service.logger.Warn("record database lock timeout requeue event failed", "taskID", task.ID, "error", eventErr)
}
w.service.logger.Warn("river async task requeued after database lock timeout",
"taskID", task.ID,
"status", queued.Status,
"riverJobID", job.ID,
)
return river.JobSnooze(time.Second)
}
}
w.service.logger.Warn("river async task completed with failure", "taskID", task.ID, "error", runErr, "riverJobID", job.ID)
return nil
}
func (s *Service) StartAsyncQueueWorker(ctx context.Context) {
if err := retryAsyncQueueStartup(
ctx,
asyncQueueStartupMaxAttempts,
asyncQueueStartupRetryDelay,
func() error { return s.startRiverQueue(ctx, true) },
func(attempt int, err error) {
s.logger.Warn("start river async queue failed; retrying",
"attempt", attempt,
"maxAttempts", asyncQueueStartupMaxAttempts,
"error", err,
)
},
); err != nil {
s.logger.Error("start river async queue failed", "error", err)
panic(err)
}
}
func (s *Service) StartAsyncQueueClient(ctx context.Context) {
if err := retryAsyncQueueStartup(
ctx,
asyncQueueStartupMaxAttempts,
asyncQueueStartupRetryDelay,
func() error { return s.startRiverQueue(ctx, false) },
func(attempt int, err error) {
s.logger.Warn("start river async queue client failed; retrying",
"attempt", attempt,
"maxAttempts", asyncQueueStartupMaxAttempts,
"error", err,
)
},
); err != nil {
s.logger.Error("start river async queue client failed", "error", err)
panic(err)
}
}
func retryAsyncQueueStartup(
ctx context.Context,
maxAttempts int,
retryDelay time.Duration,
start func() error,
onRetry func(attempt int, err error),
) error {
if maxAttempts < 1 {
maxAttempts = 1
}
var lastErr error
for attempt := 1; attempt <= maxAttempts; attempt++ {
if err := start(); err != nil {
lastErr = err
if attempt == maxAttempts {
break
}
if onRetry != nil {
onRetry(attempt, err)
}
timer := time.NewTimer(retryDelay)
select {
case <-ctx.Done():
if !timer.Stop() {
<-timer.C
}
return ctx.Err()
case <-timer.C:
}
continue
}
return nil
}
return lastErr
}
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 {
return err
}
if _, err := migrator.Migrate(ctx, rivermigrate.DirectionUp, nil); err != nil {
return err
}
controlClient, err := river.NewClient(driver, &river.Config{
ID: s.workerInstanceID + "-control",
Logger: s.logger,
TestOnly: s.cfg.AppEnv == "test",
})
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)
}
var executionClient asyncExecutionClient
if snapshot.Capacity > 0 {
executionClient, err = s.makeAsyncExecutionClient(snapshot.Capacity)
if err != nil {
return err
}
if err := executionClient.Start(ctx); err != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
_ = executionClient.StopAndCancel(cleanupCtx)
cancel()
return err
}
}
s.riverMu.Lock()
s.riverControlClient = controlClient
s.riverExecutionClient = executionClient
s.riverWorkerCapacity = snapshot.Capacity
s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
s.riverMu.Unlock()
s.observeAsyncWorkerCapacity(snapshot)
s.logger.Info("async worker capacity initialized",
"capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired,
"activeInstances", snapshot.ActiveInstances,
"globalCapacity", snapshot.GlobalCapacity,
"instanceID", snapshot.InstanceID,
"instanceHardLimit", s.cfg.AsyncWorkerInstanceHardLimit,
"hardLimit", snapshot.HardLimit,
"enabledModels", snapshot.EnabledModels,
"unlimitedModels", snapshot.UnlimitedModels,
"modelDesired", snapshot.ModelDesired,
"enabledGroups", snapshot.EnabledGroups,
"unlimitedGroups", snapshot.UnlimitedGroups,
"groupDesired", snapshot.GroupDesired,
"capped", snapshot.Capped,
)
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
if executionClient != nil {
_ = executionClient.StopAndCancel(cleanupCtx)
}
cancel()
return err
}
go s.refreshAsyncWorkerCapacity(ctx)
go s.dispatchWaitingAsyncAdmissions(ctx)
go s.reapExpiredTaskAdmissions(ctx)
go s.recoverOrphanedAsyncRiverJobs(ctx)
go s.stopAsyncWorkersOnShutdown(ctx)
return nil
}
func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.Tx], error) {
workers := river.NewWorkers()
if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
return nil, err
}
return river.NewClient(riverpgxv5.New(s.store.Pool()), &river.Config{
ID: fmt.Sprintf("%s-exec-%d-%d", s.workerInstanceID, capacity, time.Now().UnixNano()),
JobTimeout: -1,
Logger: s.logger,
CompletedJobRetentionPeriod: 24 * time.Hour,
Queues: map[string]river.QueueConfig{
asyncTaskQueueName: {MaxWorkers: capacity},
},
// Provider-backed media jobs commonly poll for 10-20 minutes. River may
// execute a still-running job again once this window elapses, so keep the
// rescue horizon above the longest configured provider poll timeout.
RescueStuckJobsAfter: time.Hour,
TestOnly: s.cfg.AppEnv == "test",
Workers: workers,
})
}
func (s *Service) makeAsyncExecutionClient(capacity int) (asyncExecutionClient, error) {
if s.asyncClientFactory != nil {
return s.asyncClientFactory(capacity)
}
return s.newRiverAsyncExecutionClient(capacity)
}
func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorkerCapacitySnapshot, error) {
if s.asyncCapacityLoader != nil {
return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
}
snapshot, err := s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
if err != nil {
return store.AsyncWorkerCapacitySnapshot{}, err
}
allocation, err := s.store.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
InstanceID: s.workerInstanceID,
PodUID: strings.TrimSpace(os.Getenv("POD_UID")),
PodName: strings.TrimSpace(os.Getenv("POD_NAME")),
Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")),
Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
DesiredCapacity: snapshot.Capacity,
CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit,
HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second,
})
if err != nil {
return store.AsyncWorkerCapacitySnapshot{}, err
}
snapshot.Capacity = allocation.Allocated
snapshot.GlobalCapacity = allocation.GlobalAllocated
snapshot.ActiveInstances = allocation.ActiveInstances
snapshot.InstanceID = allocation.InstanceID
return snapshot, nil
}
func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) {
ticker := time.NewTicker(time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
s.resizeAsyncWorkerCapacity(ctx)
}
}
}
func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
snapshot, err := s.loadAsyncWorkerCapacity(ctx)
if err != nil {
s.observeAsyncWorkerResize("refresh_failed")
s.logger.Warn("refresh async worker capacity failed; keeping current client", "error", err)
return
}
s.riverMu.RLock()
currentCapacity := s.riverWorkerCapacity
s.riverMu.RUnlock()
s.observeAsyncWorkerCapacity(snapshot)
if snapshot.Capacity == currentCapacity {
return
}
var newClient asyncExecutionClient
if snapshot.Capacity > 0 {
newClient, err = s.makeAsyncExecutionClient(snapshot.Capacity)
if err != nil {
s.observeAsyncWorkerResize("create_failed")
s.logger.Warn("create replacement async worker client failed; keeping current client",
"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
return
}
if err := newClient.Start(ctx); err != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
_ = newClient.StopAndCancel(cleanupCtx)
cancel()
s.observeAsyncWorkerResize("start_failed")
s.logger.Warn("start replacement async worker client failed; keeping current client",
"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
return
}
}
if ctx.Err() != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
if newClient != nil {
_ = newClient.StopAndCancel(cleanupCtx)
}
cancel()
return
}
s.riverMu.Lock()
oldClient := s.riverExecutionClient
s.riverExecutionClient = newClient
s.riverWorkerCapacity = snapshot.Capacity
if oldClient != nil {
if s.riverDrainingClients == nil {
s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
}
s.riverDrainingClients[oldClient] = struct{}{}
}
s.riverMu.Unlock()
s.observeAsyncWorkerResize("success")
s.logger.Info("async worker capacity resized",
"previousCapacity", currentCapacity,
"capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired,
"activeInstances", snapshot.ActiveInstances,
"globalCapacity", snapshot.GlobalCapacity,
"instanceID", snapshot.InstanceID,
"instanceHardLimit", s.cfg.AsyncWorkerInstanceHardLimit,
"hardLimit", snapshot.HardLimit,
"modelDesired", snapshot.ModelDesired,
"groupDesired", snapshot.GroupDesired,
"capped", snapshot.Capped,
)
if oldClient != nil {
go s.drainAsyncWorkerClient(oldClient)
}
}
func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {
stopCtx, cancel := context.WithTimeout(context.Background(), time.Hour)
defer cancel()
if err := client.Stop(stopCtx); err != nil {
s.logger.Warn("gracefully drain previous async worker client failed", "error", err)
}
s.riverMu.Lock()
delete(s.riverDrainingClients, client)
s.riverMu.Unlock()
}
func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
<-ctx.Done()
if s.store != nil && strings.TrimSpace(s.workerInstanceID) != "" {
drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
if err := s.store.MarkWorkerDraining(drainCtx, s.workerInstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
s.logger.Warn("mark async worker draining failed", "error", err)
}
drainCancel()
}
s.riverMu.Lock()
s.riverControlClient = nil
clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
if s.riverExecutionClient != nil {
clients = append(clients, s.riverExecutionClient)
}
for client := range s.riverDrainingClients {
clients = append(clients, client)
}
s.riverExecutionClient = nil
s.riverDrainingClients = nil
s.riverMu.Unlock()
for _, client := range clients {
stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
if err := client.StopAndCancel(stopCtx); err != nil {
s.logger.Warn("stop river async queue failed", "error", err)
}
cancel()
}
}
func (s *Service) asyncControlClient() *river.Client[pgx.Tx] {
s.riverMu.RLock()
defer s.riverMu.RUnlock()
return s.riverControlClient
}
func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacitySnapshot) {
observer, ok := s.billingMetrics.(interface {
SetAsyncWorkerCapacity(current, desired, hardLimit int, capped bool)
})
if ok {
observer.SetAsyncWorkerCapacity(snapshot.Capacity, snapshot.Desired, snapshot.HardLimit, snapshot.Capped)
}
distributedObserver, ok := s.billingMetrics.(interface {
SetDistributedWorkerCapacity(activeInstances, globalCapacity, allocatedCapacity int)
})
if ok {
distributedObserver.SetDistributedWorkerCapacity(snapshot.ActiveInstances, snapshot.GlobalCapacity, snapshot.Capacity)
}
}
func (s *Service) observeAsyncWorkerResize(outcome string) {
observer, ok := s.billingMetrics.(interface {
ObserveAsyncWorkerResize(string)
})
if ok {
observer.ObserveAsyncWorkerResize(outcome)
}
}
func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
return s.enqueueAsyncTaskWithOptions(ctx, task.ID, asyncTaskInsertOpts(task))
}
func (s *Service) enqueueAsyncTaskWithOptions(ctx context.Context, taskID string, opts *river.InsertOpts) error {
tx, err := s.store.Pool().Begin(ctx)
if err != nil {
return err
}
defer func() {
rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = tx.Rollback(rollbackCtx)
}()
if err := s.enqueueAsyncTaskTx(ctx, tx, taskID, opts); err != nil {
return err
}
return tx.Commit(ctx)
}
func (s *Service) enqueueAsyncTaskTx(ctx context.Context, tx pgx.Tx, taskID string, opts *river.InsertOpts) error {
riverClient := s.asyncControlClient()
if riverClient == nil {
return errors.New("river async queue is not started")
}
result, err := riverClient.InsertTx(ctx, tx, asyncTaskArgs{TaskID: taskID}, opts)
if err != nil {
return err
}
if result.Job != nil {
if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2,
updated_at = now()
WHERE id = $1::uuid`, taskID, result.Job.ID); err != nil {
return err
}
}
return nil
}
func (s *Service) WakeAsyncQueueAfter(ctx context.Context, delay time.Duration) {
}
func (s *Service) RunAsyncTask(ctx context.Context, task store.GatewayTask, user *auth.User) {
if err := s.EnqueueAsyncTask(ctx, task); err != nil {
s.logger.Warn("enqueue river async task failed", "taskID", task.ID, "error", err)
}
}
func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
riverClient := s.asyncControlClient()
if riverClient == nil {
return errors.New("river async queue is not started")
}
items, err := s.store.ListRecoverableAsyncTasks(ctx, 1000)
if err != nil {
return err
}
recovered := 0
for _, item := range items {
if err := s.enqueueAsyncTaskWithOptions(ctx, item.ID, asyncTaskRecoveryInsertOpts(item, time.Now())); err != nil {
return err
}
recovered++
}
if recovered > 0 {
s.logger.Info("river async queue recovered persisted tasks", "count", recovered)
}
return nil
}
func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) {
recoverJobs := func() {
runtimeRecovery, err := s.store.RecoverInterruptedRuntimeState(ctx)
if err != nil {
if ctx.Err() == nil {
s.logger.Warn("recover interrupted runtime state from worker failed", "error", err)
}
return
}
if runtimeRecovery.ReleasedConcurrencyLeases > 0 ||
runtimeRecovery.ReleasedRateReservations > 0 ||
runtimeRecovery.FailedAttempts > 0 ||
runtimeRecovery.FailedTasks > 0 ||
runtimeRecovery.RequeuedAsyncTasks > 0 ||
runtimeRecovery.CleanedTaskAdmissions > 0 {
s.logger.Warn("interrupted runtime state recovered by worker",
"releasedConcurrencyLeases", runtimeRecovery.ReleasedConcurrencyLeases,
"releasedRateReservations", runtimeRecovery.ReleasedRateReservations,
"failedAttempts", runtimeRecovery.FailedAttempts,
"failedTasks", runtimeRecovery.FailedTasks,
"requeuedAsyncTasks", runtimeRecovery.RequeuedAsyncTasks,
"cleanedTaskAdmissions", runtimeRecovery.CleanedTaskAdmissions,
)
}
recovered, err := s.store.RecoverOrphanedAsyncRiverJobs(
ctx,
orphanedRiverJobWorkerStaleAfter,
orphanedRiverJobRecoveryBatchSize,
)
if err != nil {
if ctx.Err() == nil {
s.logger.Warn("recover orphaned river jobs failed", "error", err)
}
return
}
if recovered > 0 {
s.logger.Warn("orphaned river jobs recovered", "count", recovered)
}
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
if ctx.Err() == nil {
s.logger.Warn("recover queued tasks without active river jobs failed", "error", err)
}
return
}
}
recoverJobs()
ticker := time.NewTicker(orphanedRiverJobScanInterval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
recoverJobs()
}
}
}
func asyncTaskInsertOpts(task store.GatewayTask) *river.InsertOpts {
priority := 2
if task.ID == "" {
priority = 3
}
return &river.InsertOpts{
MaxAttempts: 1000,
Priority: priority,
Queue: asyncTaskQueueName,
Tags: []string{"gateway-task"},
UniqueOpts: river.UniqueOpts{
ByArgs: true,
ByQueue: true,
ByState: []rivertype.JobState{
rivertype.JobStateAvailable,
rivertype.JobStatePending,
rivertype.JobStateRetryable,
rivertype.JobStateRunning,
rivertype.JobStateScheduled,
},
},
}
}
func asyncTaskRecoveryInsertOpts(item store.AsyncTaskQueueItem, now time.Time) *river.InsertOpts {
opts := asyncTaskInsertOpts(store.GatewayTask{ID: item.ID})
if item.NextRunAt.After(now) {
opts.ScheduledAt = item.NextRunAt
}
return opts
}
func authUserFromTask(task store.GatewayTask) *auth.User {
roles := []string{"user"}
if strings.TrimSpace(task.UserID) == "" {
roles = nil
}
return &auth.User{
ID: firstNonEmptyString(task.GatewayUserID, task.UserID),
Roles: roles,
TenantID: task.TenantID,
GatewayTenantID: task.GatewayTenantID,
TenantKey: task.TenantKey,
Source: firstNonEmptyString(task.UserSource, "gateway"),
GatewayUserID: task.GatewayUserID,
UserGroupID: task.UserGroupID,
UserGroupKey: task.UserGroupKey,
APIKeyID: task.APIKeyID,
APIKeyName: task.APIKeyName,
APIKeyPrefix: task.APIKeyPrefix,
}
}
func asyncWorkerID() string {
host, _ := os.Hostname()
host = strings.TrimSpace(host)
if host == "" {
host = "localhost"
}
return fmt.Sprintf("%s:%d:%d", host, os.Getpid(), time.Now().UnixNano())
}
var _ river.Worker[asyncTaskArgs] = (*asyncTaskWorker)(nil)