feat(routing): 引入多执行池智能调度

将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
2026-08-05 22:25:37 +08:00
parent 03c0873649
commit 7786692d32
58 changed files with 4510 additions and 348 deletions
@@ -2,6 +2,7 @@ package runner

 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
@@ -9,6 +10,7 @@ import (
 	"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"
@@ -17,7 +19,6 @@ import (
 	"github.com/riverqueue/river"
 	"github.com/riverqueue/river/riverdriver/riverpgxv5"
 	"github.com/riverqueue/river/rivermigrate"
-	"github.com/riverqueue/river/rivertype"
 )

 const (
@@ -59,6 +60,11 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
 	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))
@@ -87,6 +93,10 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
 		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)
@@ -309,14 +319,19 @@ func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.
 	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: map[string]river.QueueConfig{
-			asyncTaskQueueName: {MaxWorkers: capacity},
-		},
+		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.
@@ -346,25 +361,40 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
 		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,
-		PodUID:               strings.TrimSpace(os.Getenv("POD_UID")),
-		PodName:              strings.TrimSpace(os.Getenv("POD_NAME")),
-		Site:                 strings.TrimSpace(os.Getenv("EASYAI_SITE")),
-		Revision:             strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
-		DesiredCapacity:      snapshot.Capacity,
-		CapacityLimit:        s.cfg.AsyncWorkerInstanceHardLimit,
-		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,
+		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
@@ -599,7 +629,8 @@ func (s *Service) observeAsyncWorkerResize(outcome string) {
 }

 func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
-	return s.enqueueAsyncTaskWithOptions(ctx, task.ID, asyncTaskInsertOpts(task))
+	_, 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 {
@@ -761,22 +792,25 @@ func asyncTaskInsertOpts(task store.GatewayTask) *river.InsertOpts {
 	if task.ID == "" {
 		priority = 3
 	}
-	return &river.InsertOpts{
-		MaxAttempts: 1000,
-		Priority:    priority,
-		Queue:       asyncTaskQueueName,
-		Tags:        []string{"gateway-task"},
-		UniqueOpts: river.UniqueOpts{
-			ByArgs:  true,
-			ByQueue: true,
-			ByState: []rivertype.JobState{
-				rivertype.JobStateAvailable,
-				rivertype.JobStatePending,
-				rivertype.JobStateRetryable,
-				rivertype.JobStateRunning,
-				rivertype.JobStateScheduled,
-			},
-		},
+	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{"*"},
 	}
 }