easyai-ai-gateway/apps/api/internal/runner/queue_worker.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)
	}
	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 tx.Rollback(ctx)
	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
	}
	for _, item := range items {
		if admission, admissionErr := s.store.GetTaskAdmission(ctx, item.ID); admissionErr == nil && admission.Status == "waiting" {
			continue
		} else if admissionErr != nil && !errors.Is(admissionErr, pgx.ErrNoRows) {
			return admissionErr
		}
		if err := s.enqueueAsyncTaskWithOptions(ctx, item.ID, asyncTaskRecoveryInsertOpts(item, time.Now())); err != nil {
			return err
		}
	}
	if len(items) > 0 {
		s.logger.Info("river async queue recovered persisted tasks", "count", len(items))
	}
	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)
		}
	}
	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)