easyai-ai-gateway/apps/api/internal/runner/queue_worker.go

package runner

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/riverqueue/river"
	"github.com/riverqueue/river/riverdriver/riverpgxv5"
	"github.com/riverqueue/river/rivermigrate"
)

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
}

type workerLoadSampler interface {
	Sample(databaseConnections, databaseMax int32) workerload.ResourceSample
}

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
	}
	if assignedPool := strings.TrimSpace(task.AssignedPoolID); assignedPool != "" && assignedPool != strings.TrimSpace(w.service.cfg.ExecutionPoolID) {
		w.service.logger.Warn("river task claimed by a worker outside its assigned pool",
			"taskID", task.ID, "assignedPool", assignedPool, "workerPool", w.service.cfg.ExecutionPoolID)
		return river.JobSnooze(time.Second)
	}
	loadLease, admitted := w.service.tryStartWorkerTask()
	if !admitted {
		return river.JobSnooze(workerLoadRetryDelay(task.ID))
	}
	defer loadLease.Release()
	ctx = context.WithValue(ctx, workerLoadLeaseContextKey{}, loadLease)
	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 {
			if errors.Is(queueErr, store.ErrTaskExecutionManualReview) {
				w.service.logger.Warn("interrupted task held after upstream submission began", "taskID", task.ID, "error_category", "upstream_timeout")
				return 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.riverStore.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.recoverOrphanedAsyncRiverJobs(ctx)
	go s.stopAsyncWorkersOnShutdown(ctx)
	return nil
}

// StartAsyncAdmissionDispatcher starts the lightweight admission/reaping
// coordinator independently from River execution workers. This allows the
// dispatcher to run beside PostgreSQL while Worker pods scale on other nodes;
// the existing database task/scope locks make multiple API-site dispatchers
// safe during primary failover.
func (s *Service) StartAsyncAdmissionDispatcher(ctx context.Context) {
	s.asyncAdmissionRunner.Do(func() {
		go s.dispatchWaitingAsyncAdmissions(ctx)
		go s.reapExpiredTaskAdmissions(ctx)
		s.logger.Info("asynchronous admission dispatcher started")
	})
}

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
	}
	queues := map[string]river.QueueConfig{
		asyncTaskQueueName: {MaxWorkers: capacity},
	}
	poolQueue := executionpool.QueueName(s.cfg.ExecutionPoolID)
	if poolQueue != asyncTaskQueueName {
		queues[poolQueue] = river.QueueConfig{MaxWorkers: capacity}
	}
	return river.NewClient(riverpgxv5.New(s.riverStore.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:                      queues,
		// 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 {
		snapshot, err := s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
		if err == nil && s.workerLoad != nil {
			s.workerLoad.SetClaimLimit(snapshot.Capacity)
		}
		return snapshot, err
	}
	snapshot, err := s.coordinationStore.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
	if err != nil {
		return store.AsyncWorkerCapacitySnapshot{}, err
	}
	loadSnapshot := s.sampleWorkerLoad()
	labels := map[string]string{}
	if raw := strings.TrimSpace(s.cfg.ExecutionPoolLabels); raw != "" {
		if err := json.Unmarshal([]byte(raw), &labels); err != nil {
			return store.AsyncWorkerCapacitySnapshot{}, fmt.Errorf("decode execution pool labels: %w", err)
		}
	}
	if endpoint := strings.TrimSpace(s.cfg.WorkerAdvertiseEndpoint); endpoint != "" {
		if err := executionpool.ValidateAdvertisedEndpoint(endpoint, splitTrimmed(s.cfg.WorkerEndpointAllowedSuffixes), s.cfg.WorkerEndpointAllowPrivate); err != nil {
			return store.AsyncWorkerCapacitySnapshot{}, fmt.Errorf("validate worker advertised endpoint: %w", err)
		}
	}
	allocation, err := s.coordinationStore.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
		InstanceID:              s.workerInstanceID,
		WorkerID:                firstNonEmptyString(s.cfg.WorkerID, s.workerInstanceID),
		PoolID:                  s.cfg.ExecutionPoolID,
		Endpoint:                s.cfg.WorkerAdvertiseEndpoint,
		Labels:                  labels,
		Capabilities:            defaultWorkerCapabilities(),
		ProtocolVersion:         executionpool.ProtocolVersion,
		OrchestratorInstanceRef: s.cfg.WorkerOrchestratorInstanceRef,
		Revision:                strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
		DesiredCapacity:         snapshot.Capacity,
		CapacityLimit:           s.cfg.AsyncWorkerInstanceHardLimit,
		LoadMode:                loadSnapshot.Mode,
		SafeCapacity:            loadSnapshot.SafeCapacity,
		HeavyCapacity:           loadSnapshot.HeavyLimit,
		ActiveTasks:             loadSnapshot.ActiveTasks,
		PreparingTasks:          loadSnapshot.PreparingTasks,
		WaitingUpstreamTasks:    loadSnapshot.WaitingUpstreamTasks,
		FinalizingTasks:         loadSnapshot.FinalizingTasks,
		PressureState:           string(loadSnapshot.PressureState),
		PressureReason:          loadSnapshot.PressureReason,
		LoadSampledAt:           loadSnapshot.SampledAt,
		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
	loadSnapshot = s.workerLoad.SetClaimLimit(allocation.Allocated)
	snapshot.LoadMode = loadSnapshot.Mode
	snapshot.LocalSafeCapacity = loadSnapshot.SafeCapacity
	snapshot.LocalHeavyCapacity = loadSnapshot.HeavyLimit
	snapshot.LocalActiveTasks = loadSnapshot.ActiveTasks
	snapshot.LocalPreparingTasks = loadSnapshot.PreparingTasks
	snapshot.LocalWaitingTasks = loadSnapshot.WaitingUpstreamTasks
	snapshot.LocalFinalizingTasks = loadSnapshot.FinalizingTasks
	snapshot.LocalPressureState = string(loadSnapshot.PressureState)
	snapshot.LocalPressureReason = loadSnapshot.PressureReason
	return snapshot, nil
}

func (s *Service) sampleWorkerLoad() workerload.Snapshot {
	if s.workerLoad == nil {
		return workerload.Snapshot{Mode: workerload.ModeLegacy, SafeCapacity: s.cfg.AsyncWorkerInstanceHardLimit, HeavyLimit: s.cfg.AsyncWorkerInstanceHardLimit}
	}
	if s.workerLoadSampler == nil {
		return s.workerLoad.Snapshot()
	}
	var connections int32
	var maximum int32
	seen := make(map[*pgxpool.Pool]struct{}, 3)
	for _, database := range []*store.Store{s.store, s.coordinationStore, s.riverStore} {
		if database == nil || database.Pool() == nil {
			continue
		}
		pool := database.Pool()
		if _, ok := seen[pool]; ok {
			continue
		}
		seen[pool] = struct{}{}
		statistics := pool.Stat()
		connections += statistics.AcquiredConns()
		maximum += statistics.MaxConns()
	}
	return s.workerLoad.Observe(s.workerLoadSampler.Sample(connections, maximum))
}

func (s *Service) tryStartWorkerTask() (*workerload.Lease, bool) {
	if s.workerLoad == nil {
		return nil, true
	}
	return s.workerLoad.TryStart()
}

func workerLoadRetryDelay(taskID string) time.Duration {
	var value uint32
	for _, character := range []byte(taskID) {
		value = value*33 + uint32(character)
	}
	return 250*time.Millisecond + time.Duration(value%501)*time.Millisecond
}

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,
		"loadMode", snapshot.LoadMode,
		"localSafeCapacity", snapshot.LocalSafeCapacity,
		"localHeavyCapacity", snapshot.LocalHeavyCapacity,
		"localActiveTasks", snapshot.LocalActiveTasks,
		"pressureState", snapshot.LocalPressureState,
		"pressureReason", snapshot.LocalPressureReason,
	)
	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.coordinationStore.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)
	}
	loadObserver, ok := s.billingMetrics.(interface {
		SetWorkerLoad(activeLimit, heavyLimit, active, preparing, waiting, finalizing int, pressure string)
	})
	if ok {
		loadObserver.SetWorkerLoad(
			snapshot.LocalSafeCapacity, snapshot.LocalHeavyCapacity, snapshot.LocalActiveTasks,
			snapshot.LocalPreparingTasks, snapshot.LocalWaitingTasks, snapshot.LocalFinalizingTasks,
			snapshot.LocalPressureState,
		)
	}
}

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 {
	_, err := (riverExecutionBroker{service: s}).Publish(ctx, task.AssignedPoolID, task.ID, time.Time{})
	return err
}

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.coordinationStore.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,
			)
		}
		finalizedTerminalJobs, err := s.coordinationStore.FinalizeOrphanedTerminalAsyncRiverJobs(
			ctx,
			orphanedRiverJobWorkerStaleAfter,
			orphanedRiverJobRecoveryBatchSize,
		)
		if err != nil {
			if ctx.Err() == nil {
				s.logger.Warn("finalize orphaned terminal river jobs failed", "error", err)
			}
			return
		}
		if finalizedTerminalJobs > 0 {
			s.logger.Warn("orphaned terminal river jobs finalized", "count", finalizedTerminalJobs)
		}
		yieldedAdmissions, err := s.coordinationStore.YieldStaleAsyncTaskAdmissions(
			ctx,
			orphanedRiverJobWorkerStaleAfter,
			orphanedRiverJobRecoveryBatchSize,
		)
		if err != nil {
			if ctx.Err() == nil {
				s.logger.Warn("yield stale async task admissions failed", "error", err)
			}
			return
		}
		if yieldedAdmissions > 0 {
			s.logger.Warn("stale async task admissions yielded", "count", yieldedAdmissions)
		}
		recovered, err := s.coordinationStore.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 riverExecutionInsertOptions(task.AssignedPoolID, time.Time{}, priority)
}

func splitTrimmed(value string) []string {
	items := make([]string, 0)
	for item := range strings.SplitSeq(value, ",") {
		if item = strings.TrimSpace(item); item != "" {
			items = append(items, item)
		}
	}
	return items
}

func defaultWorkerCapabilities() map[string]any {
	return map[string]any{
		"protocolVersion": executionpool.ProtocolVersion,
		"streaming":       true,
		"media":           true,
		"taskKinds":       []any{"*"},
	}
}

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)