原因:线上 P24 验收显示 Worker 在每轮准入前逐条恢复媒体、重算候选和查询 attempts,跨地域查询导致 48 个执行槽长期只能使用约 8 至 16 个。\n\n影响:任务首次排队时保存鉴权后的 admission scope 快照,Worker 单次批量读取并只刷新动态总容量;显式重新选路保留完整慢路径。验收模拟器改到非数据库专用 Worker 节点,运行期 85% 作为立即中止硬门禁,80% 继续作为容量认证目标。\n\n验证:Go 全量测试、go vet、govulncheck、真实 PostgreSQL 迁移与跨 Store 集成测试、ShellCheck、发布及验收脚本、OpenAPI、前端 lint/test/build、pnpm audit high、Compose 与 Kubernetes 渲染均通过。
882 lines
26 KiB
Go
882 lines
26 KiB
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,
|
|
})
|
|
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)
|
|
for _, candidate := range candidates {
|
|
available, availabilityErr := s.store.RuntimeCandidateAvailable(ctx, candidate.PlatformID, candidate.PlatformModelID)
|
|
if availabilityErr != nil {
|
|
return taskAdmissionPlan{}, availabilityErr
|
|
}
|
|
if !available {
|
|
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
|
|
}
|
|
}
|
|
if !hasConcurrentLimit {
|
|
return taskAdmissionPlan{
|
|
Candidate: candidate,
|
|
Body: body,
|
|
ModelType: modelType,
|
|
}, nil
|
|
}
|
|
reservations := acceptanceInfrastructureReservations(
|
|
task,
|
|
s.rateLimitReservations(ctx, user, candidate, body),
|
|
)
|
|
if err := s.store.CheckRateLimits(ctx, reservations); err != nil {
|
|
return taskAdmissionPlan{}, err
|
|
}
|
|
return taskAdmissionPlan{
|
|
Candidate: candidate,
|
|
Body: body,
|
|
ModelType: modelType,
|
|
GroupID: groupID,
|
|
Scopes: scopes,
|
|
Eligible: true,
|
|
}, nil
|
|
}
|
|
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" && task.RunMode != "acceptance_canary" {
|
|
return scopes
|
|
}
|
|
out := append([]store.AdmissionScope(nil), scopes...)
|
|
for index := range out {
|
|
// Protocol-emulated acceptance measures Gateway and Worker capacity, so
|
|
// the isolated Run is bounded by the worker_capacity scope instead of a
|
|
// production supplier quota. The real acceptance_canary path deliberately
|
|
// retains the production platform-model concurrency limit.
|
|
if task.RunMode == "acceptance" && out[index].ScopeType == "platform_model" {
|
|
out[index].ConcurrentLimit = 0
|
|
}
|
|
if out[index].ConcurrentLimit <= 0 {
|
|
if task.RunMode != "acceptance" || out[index].ScopeType != "platform_model" {
|
|
continue
|
|
}
|
|
}
|
|
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 := make([]store.RateLimitReservation, 0, len(reservations))
|
|
for _, reservation := range reservations {
|
|
if reservation.ScopeType == "platform_model" && reservation.Metric == "concurrent" {
|
|
continue
|
|
}
|
|
out = append(out, reservation)
|
|
}
|
|
return out
|
|
}
|
|
|
|
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") ||
|
|
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
|
|
}
|
|
out := append([]store.RuntimeModelCandidate(nil), candidates...)
|
|
pinned := out[pinnedIndex]
|
|
copy(out[1:pinnedIndex+1], out[:pinnedIndex])
|
|
out[0] = pinned
|
|
return out, true
|
|
}
|
|
|
|
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
|
|
}
|
|
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
|
|
}
|
|
}
|
|
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
|
|
}
|
|
reservations := acceptanceInfrastructureReservations(
|
|
task,
|
|
s.rateLimitReservations(ctx, user, candidate, body),
|
|
)
|
|
if err := s.store.CheckRateLimits(ctx, reservations); err != nil {
|
|
return store.TaskAdmissionResult{}, true, err
|
|
}
|
|
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, error) {
|
|
if len(admissions) == 0 {
|
|
return false, nil
|
|
}
|
|
workerCapacity, err := s.coordinationStore.ActiveWorkerCapacity(ctx)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if workerCapacity <= 0 {
|
|
return true, nil
|
|
}
|
|
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 {
|
|
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) {
|
|
continue
|
|
}
|
|
return false, loadErr
|
|
}
|
|
plan, planErr := s.buildTaskAdmissionPlanForCurrentBinding(ctx, task, authUserFromTask(task), &admission)
|
|
if planErr != nil {
|
|
return false, planErr
|
|
}
|
|
if !plan.Eligible {
|
|
if enqueueErr := s.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
|
|
return false, enqueueErr
|
|
}
|
|
continue
|
|
}
|
|
input := taskAdmissionInput(task, plan, "")
|
|
if _, rebindErr := s.store.RebindWaitingTaskAdmission(ctx, input); rebindErr != nil && !errors.Is(rebindErr, pgx.ErrNoRows) {
|
|
return false, rebindErr
|
|
}
|
|
inputs = append(inputs, input)
|
|
tasksByID[task.ID] = task
|
|
}
|
|
if len(inputs) == 0 {
|
|
continue
|
|
}
|
|
outcomes, err := s.store.TryTaskAdmissionAtomicBatchWithAdmittedHook(
|
|
ctx,
|
|
inputs,
|
|
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))
|
|
},
|
|
)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
for _, outcome := range outcomes {
|
|
s.observeTaskAdmissionAttempt(outcome.Result, outcome.Err)
|
|
if outcome.Err != nil {
|
|
if errors.Is(outcome.Err, store.ErrTaskExecutionFinished) ||
|
|
errors.Is(outcome.Err, store.ErrQueueTimeout) {
|
|
continue
|
|
}
|
|
return false, outcome.Err
|
|
}
|
|
if !outcome.Result.Admitted {
|
|
return true, nil
|
|
}
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
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()
|
|
if _, changed, err := s.store.CancelQueuedTask(cleanupCtx, taskID, "client disconnected before upstream submission"); 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()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-s.asyncAdmissionWake:
|
|
case <-ticker.C:
|
|
}
|
|
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)
|
|
break
|
|
}
|
|
if len(admissions) == 0 {
|
|
break
|
|
}
|
|
saturated, err := s.dispatchWaitingAsyncTasks(ctx, admissions)
|
|
if err != nil {
|
|
s.logger.Warn("dispatch waiting async admission batch failed", "tasks", len(admissions), "error", err)
|
|
break
|
|
}
|
|
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
|
|
}
|
|
// 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:
|
|
}
|
|
result, err := s.store.ReapExpiredTaskAdmissions(ctx, 500)
|
|
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")
|
|
}
|
|
}
|
|
}
|