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

package runner

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

type taskAdmissionPlan struct {
	Candidate store.RuntimeModelCandidate
	Body      map[string]any
	ModelType string
	GroupID   string
	Scopes    []store.AdmissionScope
	Eligible  bool
}

type admissionTaskWaiter struct {
	wake     chan struct{}
	waiterID string
}

const (
	asyncWorkerCapacityScopeKey = "global"
	asyncWorkerQueueLimit       = 10000
	asyncWorkerMaxWaitSeconds   = 24 * 60 * 60
	// The dispatcher scans enough FIFO entries to fill the certified 2 x P32
	// capacity. Durable scope snapshots make this one bulk read; atomic commits
	// remain bounded by the separately configured microbatch size.
	asyncAdmissionBatchMax      = 64
	acceptanceQueueLimit        = 10000
	acceptanceQueueMaxWait      = 15 * 60
	synchronousAdmissionWakeMax = 256
)

func distributedAdmissionModelType(modelType string) bool {
	switch strings.ToLower(strings.TrimSpace(modelType)) {
	case "image_generate",
		"image_edit",
		"image_analysis",
		"image_vectorize",
		"video_generate",
		"video_enhance",
		"image_to_video",
		"text_to_video",
		"video_edit",
		"video_reference",
		"video_first_last_frame",
		"video_understanding",
		"omni_video",
		"omni",
		"audio_generate",
		"audio_understanding",
		"text_to_speech",
		"voice_clone":
		return true
	default:
		return false
	}
}

func (s *Service) buildTaskAdmissionPlan(ctx context.Context, task store.GatewayTask, user *auth.User) (taskAdmissionPlan, error) {
	return s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, user, nil)
}

func (s *Service) buildTaskAdmissionPlanForCurrentBinding(
	ctx context.Context,
	task store.GatewayTask,
	user *auth.User,
	admission *store.TaskAdmission,
) (taskAdmissionPlan, error) {
	restoredRequest, err := s.restoreTaskRequestReferences(ctx, task)
	if err != nil {
		return taskAdmissionPlan{}, err
	}
	body := normalizeRequest(task.Kind, restoredRequest)
	modelType := modelTypeFromKind(task.Kind, body)
	if !distributedAdmissionModelType(modelType) {
		return taskAdmissionPlan{Body: body, ModelType: modelType}, nil
	}
	if err := validateRequest(task.Kind, body); err != nil {
		return taskAdmissionPlan{}, parameterPreprocessClientError(err)
	}
	body, clonedVoice, err := s.resolveClonedVoiceBinding(ctx, user, task.Kind, body)
	if err != nil {
		return taskAdmissionPlan{}, err
	}
	runnerPolicy, err := s.store.GetActiveRunnerPolicy(ctx)
	if err != nil {
		return taskAdmissionPlan{}, err
	}
	cacheAffinityKeys := buildCacheAffinityKeys(task.Kind, modelType, body)
	candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user, store.ListModelCandidatesOptions{
		CacheAffinityKey:    cacheAffinityKeys.Primary,
		CacheAffinityKeys:   cacheAffinityKeys.Lookup,
		CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
		QuotaScopePrefix:    acceptanceQuotaScopePrefix(task),
	})
	if err == nil {
		candidates, err = filterCandidatesByRequestedPlatform(candidates, body)
	}
	if err == nil {
		candidates, err = filterCandidatesByClonedVoiceBinding(candidates, clonedVoice)
	}
	if err == nil {
		candidates, _, err = filterRuntimeCandidatesByRequest(task.Kind, task.Model, modelType, body, candidates)
	}
	if err == nil {
		candidates, _, err = filterRuntimeCandidatesByOutputTokens(task.Kind, task.Model, modelType, body, candidates)
	}
	if err != nil {
		return taskAdmissionPlan{}, err
	}
	candidates, _ = pinCandidatesToTaskAdmission(candidates, admission)
	var candidateRateLimitErr error
	for _, candidate := range candidates {
		available, unavailableReason, availabilityErr := s.store.RuntimeCandidateAvailability(ctx, candidate.PlatformID, candidate.PlatformModelID)
		if availabilityErr != nil {
			return taskAdmissionPlan{}, availabilityErr
		}
		if !available {
			s.observeCandidateRouting(unavailableReason)
			continue
		}
		scopes, groupID, scopeErr := s.taskAdmissionScopes(ctx, task, user, candidate)
		if scopeErr != nil {
			return taskAdmissionPlan{}, scopeErr
		}
		hasConcurrentLimit := false
		for _, scope := range scopes {
			if scope.ConcurrentLimit > 0 {
				hasConcurrentLimit = true
				break
			}
		}
		reservations := acceptanceInfrastructureReservations(
			task,
			s.rateLimitReservations(ctx, user, candidate, body),
		)
		if err := s.store.CheckRateLimits(ctx, reservations); err != nil {
			var limitErr *store.RateLimitExceededError
			if errors.As(err, &limitErr) && limitErr.ScopeType == "platform_model" {
				if candidateRateLimitErr == nil {
					candidateRateLimitErr = err
				}
				continue
			}
			return taskAdmissionPlan{}, err
		}
		if !hasConcurrentLimit {
			return taskAdmissionPlan{
				Candidate: candidate,
				Body:      body,
				ModelType: modelType,
			}, nil
		}
		return taskAdmissionPlan{
			Candidate: candidate,
			Body:      body,
			ModelType: modelType,
			GroupID:   groupID,
			Scopes:    scopes,
			Eligible:  true,
		}, nil
	}
	if candidateRateLimitErr != nil {
		s.observeCandidateRouting("all_full_queued")
		return taskAdmissionPlan{}, candidateRateLimitErr
	}
	return taskAdmissionPlan{}, store.ErrNoModelCandidate
}

func (s *Service) withAsyncWorkerCapacityScope(ctx context.Context, scopes []store.AdmissionScope) ([]store.AdmissionScope, error) {
	capacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx)
	if err != nil {
		return nil, err
	}
	if capacity <= 0 {
		return scopes, nil
	}
	out := append([]store.AdmissionScope{}, scopes...)
	out = append(out, store.AdmissionScope{
		ScopeType:       "worker_capacity",
		ScopeKey:        asyncWorkerCapacityScopeKey,
		ScopeName:       "async worker capacity",
		ConcurrentLimit: float64(capacity),
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      asyncWorkerQueueLimit,
		MaxWaitSeconds:  asyncWorkerMaxWaitSeconds,
	})
	return out, nil
}

func (s *Service) taskAdmissionScopes(
	ctx context.Context,
	task store.GatewayTask,
	user *auth.User,
	candidate store.RuntimeModelCandidate,
) ([]store.AdmissionScope, string, error) {
	scopes, groupID := s.admissionScopes(ctx, user, candidate)
	if task.AsyncMode {
		var err error
		scopes, err = s.withAsyncWorkerCapacityScope(ctx, scopes)
		if err != nil {
			return nil, "", err
		}
	}
	return acceptanceAdmissionScopes(task, scopes), groupID, nil
}

func acceptanceAdmissionScopes(task store.GatewayTask, scopes []store.AdmissionScope) []store.AdmissionScope {
	if task.RunMode != "acceptance" {
		return scopes
	}
	out := append([]store.AdmissionScope(nil), scopes...)
	for index := range out {
		out[index].ScopeKey = acceptanceScopeKey(task, out[index].ScopeKey)
		out[index].QueueLimit = acceptanceQueueLimit
		out[index].MaxWaitSeconds = acceptanceQueueMaxWait
	}
	return out
}

func acceptanceInfrastructureReservations(
	task store.GatewayTask,
	reservations []store.RateLimitReservation,
) []store.RateLimitReservation {
	if task.RunMode != "acceptance" {
		return reservations
	}
	out := append([]store.RateLimitReservation(nil), reservations...)
	for index := range out {
		out[index].ScopeKey = acceptanceScopeKey(task, out[index].ScopeKey)
	}
	return out
}

func acceptanceScopeKey(task store.GatewayTask, scopeKey string) string {
	runID := strings.TrimSpace(task.AcceptanceRunID)
	if runID == "" {
		runID = "unbound"
	}
	return "acceptance:" + runID + ":" + scopeKey
}

func acceptanceQuotaScopePrefix(task store.GatewayTask) string {
	if task.RunMode != "acceptance" {
		return ""
	}
	return acceptanceScopeKey(task, "")
}

func (s *Service) loadAsyncTaskAdmission(ctx context.Context, task store.GatewayTask) (*store.TaskAdmission, error) {
	if !task.AsyncMode {
		return nil, nil
	}
	admission, err := s.store.GetTaskAdmission(ctx, task.ID)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return &admission, nil
}

func pinCandidatesToTaskAdmission(
	candidates []store.RuntimeModelCandidate,
	admission *store.TaskAdmission,
) ([]store.RuntimeModelCandidate, bool) {
	if admission == nil ||
		(admission.Status != "waiting" && admission.Status != "admitted") ||
		!admission.ReselectRequestedAt.IsZero() ||
		len(candidates) < 2 {
		return candidates, false
	}
	pinnedIndex := -1
	for index := range candidates {
		if candidates[index].PlatformID == admission.PlatformID &&
			candidates[index].PlatformModelID == admission.PlatformModelID {
			pinnedIndex = index
			break
		}
	}
	if pinnedIndex <= 0 {
		return candidates, false
	}
	pinnedCandidate := candidates[pinnedIndex]
	if providerRateQuotaFull(pinnedCandidate) && !providerRateQuotaFull(candidates[0]) {
		return candidates, false
	}
	out := append([]store.RuntimeModelCandidate(nil), candidates...)
	pinned := out[pinnedIndex]
	copy(out[1:pinnedIndex+1], out[:pinnedIndex])
	out[0] = pinned
	return out, true
}

func providerRateQuotaFull(candidate store.RuntimeModelCandidate) bool {
	metrics := candidate.LoadMetrics
	return (metrics.RPMLimit > 0 && metrics.RPMCurrent >= metrics.RPMLimit) ||
		(metrics.TPMLimit > 0 && metrics.TPMCurrent >= metrics.TPMLimit)
}

func (s *Service) tryTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) {
	return s.tryTaskAdmissionWithAdmittedHook(ctx, task, plan, waiterID, nil)
}

func (s *Service) activeAsyncTaskAdmission(
	ctx context.Context,
	task store.GatewayTask,
	plan taskAdmissionPlan,
	admission *store.TaskAdmission,
) (store.TaskAdmissionResult, bool, error) {
	if !task.AsyncMode || admission == nil {
		return store.TaskAdmissionResult{}, false, nil
	}
	if admission.Status != "admitted" ||
		admission.PlatformID != plan.Candidate.PlatformID ||
		admission.PlatformModelID != plan.Candidate.PlatformModelID ||
		admission.UserGroupID != plan.GroupID {
		return store.TaskAdmissionResult{}, false, nil
	}
	leases, err := s.store.ActiveTaskAdmissionLeases(ctx, task.ID)
	if err != nil {
		return store.TaskAdmissionResult{}, false, err
	}
	if len(leases) == 0 {
		return store.TaskAdmissionResult{}, false, nil
	}
	leaseStore := s.coordinationStore
	if leaseStore == nil {
		leaseStore = s.store
	}
	if err := leaseStore.RenewConcurrencyLeases(ctx, leases); err != nil {
		if errors.Is(err, store.ErrConcurrencyLeaseLost) {
			return store.TaskAdmissionResult{}, false, nil
		}
		return store.TaskAdmissionResult{}, false, err
	}
	return store.TaskAdmissionResult{
		Admission: *admission,
		Admitted:  true,
		Leases:    leases,
	}, true, nil
}

func (s *Service) tryTaskAdmissionWithAdmittedHook(
	ctx context.Context,
	task store.GatewayTask,
	plan taskAdmissionPlan,
	waiterID string,
	onAdmitted func(pgx.Tx) error,
) (store.TaskAdmissionResult, error) {
	input, err := s.prepareTaskAdmissionInput(ctx, task, plan, waiterID)
	if err != nil {
		return store.TaskAdmissionResult{}, err
	}
	var result store.TaskAdmissionResult
	if onAdmitted == nil {
		result, err = s.store.TryTaskAdmission(ctx, input)
	} else {
		result, err = s.store.TryTaskAdmissionWithAdmittedHook(ctx, input, onAdmitted)
	}
	s.observeTaskAdmissionAttempt(result, err)
	return result, err
}

func (s *Service) prepareTaskAdmissionInput(
	ctx context.Context,
	task store.GatewayTask,
	plan taskAdmissionPlan,
	waiterID string,
) (store.TaskAdmissionInput, error) {
	input := taskAdmissionInput(task, plan, waiterID)
	current, currentErr := s.store.GetTaskAdmission(ctx, task.ID)
	if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) {
		return store.TaskAdmissionInput{}, currentErr
	}
	if currentErr == nil && current.Status == "waiting" &&
		(current.PlatformID != input.PlatformID || current.PlatformModelID != input.PlatformModelID || current.UserGroupID != input.UserGroupID) {
		if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil {
			// Another dispatcher may admit or remove the same FIFO row between
			// the read above and the conditional rebind. Let the atomic admission
			// transaction observe the committed state instead of aborting an
			// entire dispatcher batch on that benign race.
			if errors.Is(rebindErr, pgx.ErrNoRows) {
				return input, nil
			}
			return store.TaskAdmissionInput{}, rebindErr
		}
		s.observeTaskAdmission("candidate_migrated")
	}
	return input, nil
}

func (s *Service) observeTaskAdmissionAttempt(result store.TaskAdmissionResult, err error) {
	if result.NewlyAdmitted {
		s.observeTaskAdmission("admitted")
		s.observeTaskAdmissionWait(time.Since(result.Admission.EnqueuedAt))
	}
	var limitErr *store.RateLimitExceededError
	switch {
	case errors.As(err, &limitErr) && limitErr.Reason == "queue_full":
		s.observeTaskAdmission("queue_full")
	case errors.Is(err, store.ErrQueueTimeout):
		s.observeTaskAdmission("timeout")
	}
}

func taskAdmissionInput(task store.GatewayTask, plan taskAdmissionPlan, waiterID string) store.TaskAdmissionInput {
	mode := "sync"
	if task.AsyncMode {
		mode = "async"
	}
	return store.TaskAdmissionInput{
		TaskID:          task.ID,
		PlatformID:      plan.Candidate.PlatformID,
		PlatformModelID: plan.Candidate.PlatformModelID,
		UserGroupID:     plan.GroupID,
		QueueKey:        plan.Candidate.QueueKey,
		Mode:            mode,
		Priority:        task.Priority,
		WaiterID:        waiterID,
		Scopes:          plan.Scopes,
	}
}

func (s *Service) waitForSynchronousAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan) (store.TaskAdmissionResult, error) {
	waiterID := uuid.NewString()
	return s.waitForTaskAdmission(ctx, task, plan, waiterID)
}

func (s *Service) waitForTaskAdmission(ctx context.Context, task store.GatewayTask, plan taskAdmissionPlan, waiterID string) (store.TaskAdmissionResult, error) {
	taskWake, unregister := s.registerAdmissionTaskWaiter(task.ID, waiterID)
	defer unregister()
	pollTimer := time.NewTimer(admissionPollInterval(task.ID))
	defer pollTimer.Stop()
	for {
		result, err := s.tryTaskAdmission(ctx, task, plan, waiterID)
		if err != nil || result.Admitted {
			return result, err
		}
		select {
		case <-ctx.Done():
			_ = s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
			s.observeTaskAdmission("cancelled")
			return store.TaskAdmissionResult{}, ctx.Err()
		case <-taskWake:
		case <-pollTimer.C:
			pollTimer.Reset(admissionPollInterval(task.ID))
		}
	}
}

func admissionPollInterval(taskID string) time.Duration {
	hasher := fnv.New32a()
	_, _ = hasher.Write([]byte(taskID))
	return 30*time.Second + time.Duration(hasher.Sum32()%5000)*time.Millisecond
}

func (s *Service) ensureCandidateAdmission(
	ctx context.Context,
	task store.GatewayTask,
	user *auth.User,
	body map[string]any,
	candidate store.RuntimeModelCandidate,
) (store.TaskAdmissionResult, bool, error) {
	scopes, groupID, err := s.taskAdmissionScopes(ctx, task, user, candidate)
	if err != nil {
		return store.TaskAdmissionResult{}, true, err
	}
	hasConcurrentLimit := false
	for _, scope := range scopes {
		if scope.ConcurrentLimit > 0 {
			hasConcurrentLimit = true
			break
		}
	}
	reservations := acceptanceInfrastructureReservations(
		task,
		s.rateLimitReservations(ctx, user, candidate, body),
	)
	if err := s.store.CheckRateLimits(ctx, reservations); err != nil {
		return store.TaskAdmissionResult{}, true, err
	}
	if !hasConcurrentLimit {
		if err := s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
			return store.TaskAdmissionResult{}, false, err
		}
		return store.TaskAdmissionResult{}, false, nil
	}
	plan := taskAdmissionPlan{
		Candidate: candidate,
		Body:      body,
		ModelType: candidate.ModelType,
		GroupID:   groupID,
		Scopes:    scopes,
		Eligible:  true,
	}
	waiterID := ""
	if !task.AsyncMode {
		waiterID = uuid.NewString()
	}
	result, err := s.waitForTaskAdmission(ctx, task, plan, waiterID)
	return result, true, err
}

func (s *Service) observeTaskAdmission(event string) {
	observer, ok := s.billingMetrics.(interface {
		ObserveTaskAdmission(string)
	})
	if ok {
		observer.ObserveTaskAdmission(event)
	}
}

func (s *Service) observeTaskAdmissionWait(wait time.Duration) {
	observer, ok := s.billingMetrics.(interface {
		ObserveTaskAdmissionWait(time.Duration)
	})
	if ok {
		observer.ObserveTaskAdmissionWait(wait)
	}
}

func (s *Service) StartAdmissionNotifier(ctx context.Context) {
	s.admissionListener.Do(func() {
		go s.dispatchWaitingSynchronousAdmissions(ctx)
		go s.renewSynchronousAdmissionWaiters(ctx)
		go func() {
			for ctx.Err() == nil {
				err := s.store.ListenTaskAdmissionNotifications(ctx, func(taskID string) {
					s.signalAdmissionWake(taskID)
				})
				if ctx.Err() != nil {
					return
				}
				if s.logger != nil {
					s.logger.Warn("task admission LISTEN connection interrupted; periodic polling remains active", "error", err)
				}
				timer := time.NewTimer(time.Second)
				select {
				case <-ctx.Done():
					timer.Stop()
					return
				case <-timer.C:
				}
			}
		}()
	})
}

func (s *Service) dispatchWaitingSynchronousAdmissions(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-s.admissionWake:
		}
		taskIDs, err := s.store.ListWaitingTaskAdmissionIDs(ctx, synchronousAdmissionWakeMax)
		if err != nil {
			if s.logger != nil {
				s.logger.Warn("list waiting synchronous admissions failed", "error", err)
			}
			continue
		}
		s.admissionWakeMu.Lock()
		waiters := make([]chan struct{}, 0, len(taskIDs))
		for _, taskID := range taskIDs {
			if waiter := s.admissionTaskWaiters[taskID]; waiter != nil {
				waiters = append(waiters, waiter.wake)
			}
		}
		s.admissionWakeMu.Unlock()
		for _, taskWake := range waiters {
			select {
			case taskWake <- struct{}{}:
			default:
			}
		}
	}
}

func (s *Service) renewSynchronousAdmissionWaiters(ctx context.Context) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		s.admissionWakeMu.Lock()
		taskIDs := make([]string, 0, len(s.admissionTaskWaiters))
		waiterIDs := make([]string, 0, len(s.admissionTaskWaiters))
		for taskID, waiter := range s.admissionTaskWaiters {
			if waiter == nil || strings.TrimSpace(waiter.waiterID) == "" {
				continue
			}
			taskIDs = append(taskIDs, taskID)
			waiterIDs = append(waiterIDs, waiter.waiterID)
		}
		s.admissionWakeMu.Unlock()
		if err := s.store.RenewTaskAdmissionWaiters(ctx, taskIDs, waiterIDs); err != nil && s.logger != nil {
			s.logger.Warn("renew synchronous admission waiters failed", "waiterCount", len(taskIDs), "error", err)
		}
	}
}

func (s *Service) registerAdmissionTaskWaiter(taskID string, waiterID string) (<-chan struct{}, func()) {
	s.admissionWakeMu.Lock()
	waiter := &admissionTaskWaiter{wake: make(chan struct{}, 1), waiterID: waiterID}
	s.admissionTaskWaiters[taskID] = waiter
	s.admissionWakeMu.Unlock()
	return waiter.wake, func() {
		s.admissionWakeMu.Lock()
		if s.admissionTaskWaiters[taskID] == waiter {
			delete(s.admissionTaskWaiters, taskID)
		}
		s.admissionWakeMu.Unlock()
	}
}

func (s *Service) signalAdmissionWake(taskID string) {
	taskID = strings.TrimSpace(taskID)
	if taskID == "" || taskID == "*" {
		select {
		case s.admissionWake <- struct{}{}:
		default:
		}
	} else {
		s.admissionWakeMu.Lock()
		waiter := s.admissionTaskWaiters[taskID]
		s.admissionWakeMu.Unlock()
		if waiter != nil {
			select {
			case waiter.wake <- struct{}{}:
			default:
			}
		}
	}
	select {
	case s.asyncAdmissionWake <- struct{}{}:
	default:
	}
}

func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) error {
	if !task.AsyncMode {
		return errors.New("only asynchronous tasks can be submitted to the async queue")
	}
	if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
		return ctx.Err()
	}
	user := authUserFromTask(task)
	plan, err := s.buildTaskAdmissionPlan(ctx, task, user)
	if err != nil {
		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
			return ctx.Err()
		}
		if store.ModelCandidateRetryAfter(err) > 0 ||
			(errors.Is(err, store.ErrRateLimited) && store.RateLimitRetryable(err)) {
			if enqueueErr := s.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
				_, _ = s.store.FailQueuedTask(
					context.WithoutCancel(ctx),
					task.ID,
					"enqueue_failed",
					enqueueErr.Error(),
				)
				return enqueueErr
			}
			return nil
		}
		_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
		return err
	}
	if !plan.Eligible {
		if err := s.EnqueueAsyncTask(ctx, task); err != nil {
			if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
				return ctx.Err()
			}
			_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, "enqueue_failed", err.Error())
			return err
		}
		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
			return ctx.Err()
		}
		return nil
	}
	input := taskAdmissionInput(task, plan, "")
	current, currentErr := s.store.GetTaskAdmission(ctx, task.ID)
	if currentErr != nil && !errors.Is(currentErr, pgx.ErrNoRows) {
		return currentErr
	}
	if currentErr == nil && current.Status == "waiting" &&
		(current.PlatformID != input.PlatformID || current.PlatformModelID != input.PlatformModelID || current.UserGroupID != input.UserGroupID) {
		if _, err := s.store.RebindWaitingTaskAdmission(ctx, input); err != nil {
			return err
		}
		s.observeTaskAdmission("candidate_migrated")
	}
	// Register FIFO order without creating a River job. A Worker dispatcher
	// first reserves both business concurrency and a live execution slot, then
	// creates the unique River job in the same transaction.
	if _, err := s.store.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
		if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
			return ctx.Err()
		}
		_, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error())
		return err
	}
	if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) {
		return ctx.Err()
	}
	return nil
}

func admissionInputFromSnapshot(admission store.TaskAdmission, workerCapacity int) store.TaskAdmissionInput {
	scopes := append([]store.AdmissionScope(nil), admission.Scopes...)
	for index := range scopes {
		if scopes[index].ScopeType == "worker_capacity" && scopes[index].ScopeKey == asyncWorkerCapacityScopeKey {
			scopes[index].ConcurrentLimit = float64(workerCapacity)
		}
	}
	return store.TaskAdmissionInput{
		TaskID:          admission.TaskID,
		PlatformID:      admission.PlatformID,
		PlatformModelID: admission.PlatformModelID,
		UserGroupID:     admission.UserGroupID,
		QueueKey:        admission.QueueKey,
		Mode:            admission.Mode,
		Priority:        admission.Priority,
		Scopes:          scopes,
	}
}

func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []store.TaskAdmission) (bool, bool, error) {
	if len(admissions) == 0 {
		return false, false, nil
	}
	workerCapacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx)
	if err != nil {
		return false, false, err
	}
	if workerCapacity <= 0 {
		s.observeAsyncAdmissionDispatch("global_wait")
		return true, false, nil
	}
	progressed := false
	microbatchSize := s.cfg.AsyncAdmissionMicrobatchSize
	if microbatchSize <= 0 {
		microbatchSize = 1
	}
	for start := 0; start < len(admissions); start += microbatchSize {
		end := min(start+microbatchSize, len(admissions))
		inputs := make([]store.TaskAdmissionInput, 0, end-start)
		tasksByID := make(map[string]store.GatewayTask, end-start)
		for _, admission := range admissions[start:end] {
			if admission.ReselectRequestedAt.IsZero() && len(admission.Scopes) > 0 && !admission.SnapshotInvalid {
				inputs = append(inputs, admissionInputFromSnapshot(admission, workerCapacity))
				tasksByID[admission.TaskID] = store.GatewayTask{ID: admission.TaskID}
				continue
			}
			task, loadErr := s.store.GetTask(ctx, admission.TaskID)
			if loadErr != nil {
				if errors.Is(loadErr, pgx.ErrNoRows) {
					progressed = true
					continue
				}
				return false, progressed, loadErr
			}
			plan, planErr := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, authUserFromTask(task), &admission)
			if planErr != nil {
				handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, task, admission, planErr)
				if handleErr != nil {
					return false, progressed, handleErr
				}
				progressed = progressed || handled
				continue
			}
			if !plan.Eligible {
				// This task already owns a durable admission row from its previous
				// binding. Promote that row and insert the River job in the same
				// transaction even when the replacement candidate has no business
				// concurrency scope; direct enqueue would strand a waiting row and
				// make every dispatcher enqueue it forever.
				if plan.Candidate.PlatformID == "" || plan.Candidate.PlatformModelID == "" {
					handled, handleErr := s.handleAsyncAdmissionTaskError(
						ctx,
						task,
						admission,
						errors.New("waiting admission rebuilt without a platform model candidate"),
					)
					if handleErr != nil {
						return false, progressed, handleErr
					}
					progressed = progressed || handled
					continue
				}
			}
			input := taskAdmissionInput(task, plan, "")
			if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil && !errors.Is(rebindErr, pgx.ErrNoRows) {
				handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, task, admission, rebindErr)
				if handleErr != nil {
					return false, progressed, handleErr
				}
				progressed = progressed || handled
				continue
			}
			progressed = true
			inputs = append(inputs, input)
			tasksByID[task.ID] = task
		}
		if len(inputs) == 0 {
			continue
		}
		enqueueHook := func(tx pgx.Tx, input store.TaskAdmissionInput) error {
			task, ok := tasksByID[input.TaskID]
			if !ok {
				return fmt.Errorf("async admission batch task %s was not prepared", input.TaskID)
			}
			return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task))
		}
		outcomes, batchErr := s.store.TryTaskAdmissionAtomicBatchWithAdmittedHook(
			ctx,
			inputs,
			enqueueHook,
		)
		if batchErr != nil {
			if isAsyncAdmissionSystemError(batchErr) {
				return false, progressed, batchErr
			}
			failedTaskID := asyncAdmissionFailedTaskID(outcomes)
			if failedTaskID == "" {
				return false, progressed, batchErr
			}
			failedTask, loadErr := s.asyncAdmissionDispatchTask(ctx, tasksByID[failedTaskID])
			if loadErr != nil {
				return false, progressed, loadErr
			}
			failedAdmission := admissionByTaskID(admissions[start:end], failedTaskID)
			handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, failedTask, failedAdmission, batchErr)
			if handleErr != nil {
				return false, progressed, handleErr
			}
			progressed = progressed || handled

			for _, input := range inputs {
				if input.TaskID == failedTaskID {
					continue
				}
				result, oneErr := s.store.TryTaskAdmissionWithAdmittedHook(ctx, input, func(tx pgx.Tx) error {
					return enqueueHook(tx, input)
				})
				globalSaturated, outcomeProgress, outcomeErr := s.handleAsyncAdmissionOutcome(
					ctx,
					store.TaskAdmissionBatchOutcome{TaskID: input.TaskID, Result: result, Err: oneErr},
					input,
					tasksByID,
					admissionByTaskID(admissions[start:end], input.TaskID),
				)
				if outcomeErr != nil {
					return false, progressed, outcomeErr
				}
				progressed = progressed || outcomeProgress
				if globalSaturated {
					return true, progressed, nil
				}
			}
			continue
		}
		for _, outcome := range outcomes {
			input := inputByTaskID(inputs, outcome.TaskID)
			globalSaturated, outcomeProgress, outcomeErr := s.handleAsyncAdmissionOutcome(
				ctx,
				outcome,
				input,
				tasksByID,
				admissionByTaskID(admissions[start:end], outcome.TaskID),
			)
			if outcomeErr != nil {
				return false, progressed, outcomeErr
			}
			progressed = progressed || outcomeProgress
			if globalSaturated {
				return true, progressed, nil
			}
		}
	}
	return false, progressed, nil
}

func (s *Service) handleAsyncAdmissionOutcome(
	ctx context.Context,
	outcome store.TaskAdmissionBatchOutcome,
	input store.TaskAdmissionInput,
	tasksByID map[string]store.GatewayTask,
	admission store.TaskAdmission,
) (bool, bool, error) {
	s.observeTaskAdmissionAttempt(outcome.Result, outcome.Err)
	if outcome.Err != nil {
		if errors.Is(outcome.Err, store.ErrTaskExecutionFinished) || errors.Is(outcome.Err, store.ErrQueueTimeout) {
			return false, true, nil
		}
		if isAsyncAdmissionSystemError(outcome.Err) {
			return false, false, outcome.Err
		}
		task, err := s.asyncAdmissionDispatchTask(ctx, tasksByID[outcome.TaskID])
		if err != nil {
			return false, false, err
		}
		handled, handleErr := s.handleAsyncAdmissionTaskError(ctx, task, admission, outcome.Err)
		return false, handled, handleErr
	}
	if outcome.Result.Admitted {
		s.observeAsyncAdmissionDispatch("admitted")
		return false, true, nil
	}

	globalBlocked := false
	groupBlocked := false
	platformSaturated := false
	for _, blocker := range outcome.Result.Blockers {
		switch blocker.ScopeType {
		case "worker_capacity":
			if blocker.Reason == "saturated" {
				globalBlocked = true
			}
		case "user_group":
			groupBlocked = true
		case "platform_model":
			if blocker.Reason == "saturated" {
				platformSaturated = true
			}
		}
	}
	if globalBlocked {
		s.observeAsyncAdmissionDispatch("global_wait")
		return true, false, nil
	}
	if groupBlocked {
		s.observeAsyncAdmissionDispatch("group_wait")
		return false, false, nil
	}
	if platformSaturated {
		platformModelID := outcome.Result.Admission.PlatformModelID
		if platformModelID == "" {
			platformModelID = input.PlatformModelID
		}
		_, marked, err := s.store.RequestTaskAdmissionReselect(ctx, outcome.TaskID, platformModelID)
		if err != nil {
			return false, false, err
		}
		s.observeAsyncAdmissionDispatch("platform_wait")
		if marked {
			s.observeTaskAdmission("candidate_reselect_requested")
			s.observeAsyncAdmissionDispatch("reselect")
			return false, true, nil
		}
		return false, false, nil
	}
	if len(outcome.Result.Blockers) > 0 {
		s.observeAsyncAdmissionDispatch("platform_wait")
		return false, false, nil
	}

	task, err := s.asyncAdmissionDispatchTask(ctx, tasksByID[outcome.TaskID])
	if err != nil {
		return false, false, err
	}
	handled, handleErr := s.handleAsyncAdmissionTaskError(
		ctx,
		task,
		admission,
		errors.New("task admission returned neither admitted nor a blocker"),
	)
	return false, handled, handleErr
}

func (s *Service) handleAsyncAdmissionTaskError(
	ctx context.Context,
	task store.GatewayTask,
	admission store.TaskAdmission,
	cause error,
) (bool, error) {
	if isAsyncAdmissionSystemError(cause) {
		return false, cause
	}
	now := time.Now()
	terminal := asyncAdmissionTaskErrorTerminal(cause)
	code := asyncAdmissionTaskErrorCode(cause)
	if !admission.WaitDeadlineAt.IsZero() && !admission.WaitDeadlineAt.After(now) {
		terminal = true
		if !errors.Is(cause, store.ErrNoModelCandidate) {
			code = "admission_dispatch_failed"
		}
	}
	if terminal {
		callbackURL, callbackErr := s.taskCallbackURL(ctx, task.ID)
		if callbackErr != nil {
			return false, callbackErr
		}
		_, failErr := s.store.FailQueuedTaskWithCallback(
			context.WithoutCancel(ctx),
			task.ID,
			code,
			cause.Error(),
			callbackURL,
			task.RunMode == "simulation",
		)
		if errors.Is(failErr, store.ErrTaskExecutionFinished) {
			return true, nil
		}
		if failErr != nil {
			return false, failErr
		}
		s.observeAsyncAdmissionDispatch("terminal")
		if s.logger != nil {
			s.logger.Info("terminally failed waiting async admission",
				"taskID", task.ID,
				"error_category", code,
			)
		}
		return true, nil
	}

	retryAt := asyncAdmissionRetryAt(now, task.ID, admission.DispatchFailureCount, cause)
	updated, changed, deferErr := s.store.DeferWaitingTaskAdmission(
		ctx,
		task.ID,
		admission.PlatformModelID,
		retryAt,
		"deferred",
		code,
	)
	if deferErr != nil {
		return false, deferErr
	}
	if !changed {
		return true, nil
	}
	s.observeAsyncAdmissionDispatch("deferred")
	s.observeAsyncAdmissionDispatch("task_error")
	if s.logger != nil && (updated.DispatchFailureCount == 1 || updated.DispatchFailureCount == 5 || updated.DispatchFailureCount == 7 || updated.DispatchFailureCount%10 == 0) {
		s.logger.Warn("deferred task-specific async admission failure",
			"taskID", task.ID,
			"failureCount", updated.DispatchFailureCount,
			"nextDispatchAt", updated.DispatchNextAt.UTC().Format(time.RFC3339),
			"error_category", code,
		)
	}
	return true, nil
}

func asyncAdmissionTaskErrorTerminal(err error) bool {
	if errors.Is(err, store.ErrNoModelCandidate) {
		return store.ModelCandidateRetryAfter(err) <= 0 && store.ModelCandidateRecoveryAt(err).IsZero()
	}
	if errors.Is(err, store.ErrInvalidPlatformModelConfiguration) {
		return true
	}
	if errors.Is(err, store.ErrRateLimited) {
		return !store.RateLimitRetryable(err)
	}
	var clientErr *clients.ClientError
	return errors.As(err, &clientErr) && !clientErr.Retryable
}

func asyncAdmissionTaskErrorCode(err error) string {
	if errors.Is(err, store.ErrNoModelCandidate) {
		return store.ModelCandidateErrorCode(err)
	}
	return clients.ErrorCode(err)
}

func asyncAdmissionRetryAt(now time.Time, taskID string, failureCount int, err error) time.Time {
	if recoveryAt := store.ModelCandidateRecoveryAt(err); recoveryAt.After(now) {
		return recoveryAt
	}
	if delay := store.ModelCandidateRetryAfter(err); delay > 0 {
		return now.Add(delay)
	}
	if delay := store.RateLimitRetryAfter(err); delay > 0 {
		return now.Add(delay)
	}
	delay := time.Second
	for index := 0; index < failureCount && delay < time.Minute; index++ {
		delay *= 2
	}
	if delay > time.Minute {
		delay = time.Minute
	}
	hasher := fnv.New32a()
	_, _ = hasher.Write([]byte(taskID))
	jitterPermille := int64(hasher.Sum32()%401) - 200
	delay += time.Duration(int64(delay) * jitterPermille / 1000)
	if delay < time.Second {
		delay = time.Second
	}
	if delay > time.Minute {
		delay = time.Minute
	}
	return now.Add(delay)
}

func isAsyncAdmissionSystemError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || store.IsPostgresUnavailable(err) {
		return true
	}
	var sqlState interface{ SQLState() string }
	if !errors.As(err, &sqlState) {
		return false
	}
	code := strings.TrimSpace(sqlState.SQLState())
	return strings.HasPrefix(code, "08") || // connection exception
		strings.HasPrefix(code, "25") || // invalid transaction state
		strings.HasPrefix(code, "40") || // transaction rollback/deadlock
		strings.HasPrefix(code, "53") || // insufficient resources
		strings.HasPrefix(code, "58") || // system error
		code == "55P03" || // lock not available
		code == "57014" || // query cancelled/statement timeout
		code == "57P01" || code == "57P02" || code == "57P03"
}

func asyncAdmissionFailedTaskID(outcomes []store.TaskAdmissionBatchOutcome) string {
	for index := len(outcomes) - 1; index >= 0; index-- {
		if outcomes[index].Err != nil {
			return outcomes[index].TaskID
		}
	}
	return ""
}

func inputByTaskID(inputs []store.TaskAdmissionInput, taskID string) store.TaskAdmissionInput {
	for _, input := range inputs {
		if input.TaskID == taskID {
			return input
		}
	}
	return store.TaskAdmissionInput{TaskID: taskID}
}

func admissionByTaskID(admissions []store.TaskAdmission, taskID string) store.TaskAdmission {
	for _, admission := range admissions {
		if admission.TaskID == taskID {
			return admission
		}
	}
	return store.TaskAdmission{TaskID: taskID}
}

func (s *Service) asyncAdmissionDispatchTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
	if task.ID == "" {
		return store.GatewayTask{}, pgx.ErrNoRows
	}
	if task.Kind != "" {
		return task, nil
	}
	return s.store.GetTask(ctx, task.ID)
}

func (s *Service) observeAsyncAdmissionDispatch(outcome string) {
	observer, ok := s.billingMetrics.(interface {
		ObserveAsyncAdmissionDispatch(string)
	})
	if ok {
		observer.ObserveAsyncAdmissionDispatch(outcome)
	}
}

func (s *Service) cancelAsyncSubmissionIfDisconnected(ctx context.Context, taskID string) bool {
	if ctx.Err() == nil {
		return false
	}
	cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
	defer cancel()
	callbackURL, callbackErr := s.taskCallbackURL(cleanupCtx, taskID)
	if callbackErr != nil && s.logger != nil {
		s.logger.Warn("resolve disconnected task callback failed", "taskID", taskID, "error", callbackErr)
	}
	if _, changed, err := s.store.CancelQueuedTaskWithCallback(
		cleanupCtx,
		taskID,
		"client disconnected before upstream submission",
		callbackURL,
	); err != nil {
		if s.logger != nil {
			s.logger.Warn("cancel disconnected asynchronous task failed", "taskID", taskID, "error", err)
		}
	} else if changed {
		s.observeTaskAdmission("cancelled")
	}
	return true
}

func (s *Service) dispatchWaitingAsyncAdmissions(ctx context.Context) {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	systemBackoff := time.Second
	var systemRetryAt time.Time
	for {
		select {
		case <-ctx.Done():
			return
		case <-s.asyncAdmissionWake:
		case <-ticker.C:
		}
		if time.Now().Before(systemRetryAt) {
			continue
		}
		for {
			batchLimit := asyncAdmissionBatchLimit(s.cfg.AsyncWorkerHardLimit)
			admissions, err := s.store.ListWaitingAsyncAdmissions(ctx, batchLimit)
			if err != nil {
				s.logger.Warn("list waiting async admissions failed", "error", err)
				s.observeAsyncAdmissionDispatch("system_error")
				systemRetryAt = time.Now().Add(systemBackoff)
				systemBackoff = min(systemBackoff*2, 30*time.Second)
				break
			}
			if len(admissions) == 0 {
				systemBackoff = time.Second
				systemRetryAt = time.Time{}
				break
			}
			saturated, progressed, err := s.dispatchWaitingAsyncTasks(ctx, admissions)
			if err != nil {
				s.logger.Warn("dispatch waiting async admission batch failed", "tasks", len(admissions), "error", err)
				s.observeAsyncAdmissionDispatch("system_error")
				systemRetryAt = time.Now().Add(systemBackoff)
				systemBackoff = min(systemBackoff*2, 30*time.Second)
				break
			}
			systemBackoff = time.Second
			systemRetryAt = time.Time{}
			if saturated {
				// Every asynchronous task shares the global worker-capacity FIFO
				// scope. Once its current head cannot be admitted, later tasks
				// cannot make progress until a lease is released.
				s.observeTaskAdmission("dispatch_saturated")
				break
			}
			if !progressed {
				break
			}
			// PostgreSQL notifications are wake-up hints and collapse while a
			// batch is being admitted. Keep filling consecutive batches until
			// the global execution scope is actually saturated so a large burst
			// cannot strand most Worker slots behind the one-second fallback.
		}
	}
}

func asyncAdmissionBatchLimit(globalHardLimit int) int {
	if globalHardLimit <= 0 || globalHardLimit > asyncAdmissionBatchMax {
		return asyncAdmissionBatchMax
	}
	return globalHardLimit
}

func (s *Service) reapExpiredTaskAdmissions(ctx context.Context) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		defaultCallbackURL := ""
		if s.cfg.TaskProgressCallbackEnabled {
			defaultCallbackURL = s.cfg.TaskProgressCallbackURL
		}
		result, err := s.store.ReapExpiredTaskAdmissions(ctx, 500, defaultCallbackURL)
		if err != nil {
			if s.logger != nil {
				s.logger.Warn("reap expired task admissions failed", "error", err)
			}
			continue
		}
		for range result.ExpiredWaiters {
			s.observeTaskAdmission("expired")
		}
		for range result.ExpiredDeadlines {
			s.observeTaskAdmission("timeout")
		}
	}
}