BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。

实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。

新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
This commit is contained in:
wangbo
2026-08-05 22:25:37 +08:00
parent 03c0873649
commit 7786692d32
58 changed files with 4510 additions and 348 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/worker_registry.go
+143 -51
View File
@@ -2,6 +2,7 @@ package store
import (
"context"
"encoding/json"
"errors"
"strings"
"time"
@@ -9,31 +10,37 @@ import (
"github.com/jackc/pgx/v5"
)
const workerHeartbeatStaleAfter = 30 * time.Second
const workerHeartbeatStaleAfter = 15 * time.Second
type WorkerRegistrationInput struct {
InstanceID string
PodUID string
PodName string
Site string
Revision string
DesiredCapacity int
CapacityLimit int
LoadMode string
SafeCapacity int
HeavyCapacity int
ActiveTasks int
PreparingTasks int
WaitingUpstreamTasks int
FinalizingTasks int
PressureState string
PressureReason string
LoadSampledAt time.Time
HeartbeatStaleAfter time.Duration
InstanceID string
WorkerID string
PoolID string
Endpoint string
Labels map[string]string
Capabilities map[string]any
ProtocolVersion string
OrchestratorInstanceRef string
Revision string
DesiredCapacity int
CapacityLimit int
LoadMode string
SafeCapacity int
HeavyCapacity int
ActiveTasks int
PreparingTasks int
WaitingUpstreamTasks int
FinalizingTasks int
PressureState string
PressureReason string
LoadSampledAt time.Time
HeartbeatStaleAfter time.Duration
}
type WorkerAllocation struct {
InstanceID string
WorkerID string
PoolID string
DesiredCapacity int
Allocated int
GlobalAllocated int
@@ -63,6 +70,31 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
if input.InstanceID == "" {
return WorkerAllocation{}, errors.New("worker instance ID is required")
}
input.WorkerID = strings.TrimSpace(input.WorkerID)
if input.WorkerID == "" {
input.WorkerID = input.InstanceID
}
input.PoolID = strings.TrimSpace(input.PoolID)
if input.PoolID == "" {
input.PoolID = "legacy-default"
}
if strings.TrimSpace(input.ProtocolVersion) == "" {
input.ProtocolVersion = "v1"
}
if input.Labels == nil {
input.Labels = map[string]string{}
}
if input.Capabilities == nil {
input.Capabilities = map[string]any{}
}
labels, err := json.Marshal(input.Labels)
if err != nil {
return WorkerAllocation{}, err
}
capabilities, err := json.Marshal(input.Capabilities)
if err != nil {
return WorkerAllocation{}, err
}
if input.DesiredCapacity < 0 {
return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
}
@@ -110,22 +142,38 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
return WorkerAllocation{}, err
}
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_execution_pools (pool_id, labels, capabilities, state, updated_at)
VALUES ($1, $2::jsonb, $3::jsonb, 'active', now())
ON CONFLICT (pool_id) DO UPDATE
SET labels = EXCLUDED.labels,
capabilities = EXCLUDED.capabilities,
updated_at = now()
WHERE gateway_execution_pools.state <> 'disabled'`, input.PoolID, labels, capabilities); err != nil {
return WorkerAllocation{}, err
}
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_worker_instances (
instance_id, pod_uid, pod_name, site, revision, status,
instance_id, worker_id, pool_id, endpoint, labels, capabilities, protocol_version,
orchestrator_instance_ref, revision, status,
desired_capacity, capacity_limit, hard_capacity_limit, safe_capacity, heavy_capacity,
active_tasks, preparing_tasks, waiting_upstream_tasks, finalizing_tasks,
pressure_state, pressure_reason, load_sampled_at,
allocated_capacity, started_at, heartbeat_at, updated_at
)
VALUES (
$1, $2, $3, $4, $5, 'active', $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17,
$1, $2, $3, $4, $5::jsonb, $6::jsonb, $7,
$8, $9, 'active', $10, $11, $12, $13, $14,
$15, $16, $17, $18, $19, $20, $21,
0, now(), now(), now()
)
ON CONFLICT (instance_id) DO UPDATE
SET pod_uid = EXCLUDED.pod_uid,
pod_name = EXCLUDED.pod_name,
site = EXCLUDED.site,
SET worker_id = EXCLUDED.worker_id,
pool_id = EXCLUDED.pool_id,
endpoint = EXCLUDED.endpoint,
labels = EXCLUDED.labels,
capabilities = EXCLUDED.capabilities,
protocol_version = EXCLUDED.protocol_version,
orchestrator_instance_ref = EXCLUDED.orchestrator_instance_ref,
revision = EXCLUDED.revision,
status = CASE
WHEN gateway_worker_instances.status = 'draining' THEN 'draining'
@@ -146,9 +194,13 @@ SET pod_uid = EXCLUDED.pod_uid,
heartbeat_at = now(),
updated_at = now()`,
input.InstanceID,
strings.TrimSpace(input.PodUID),
strings.TrimSpace(input.PodName),
strings.TrimSpace(input.Site),
input.WorkerID,
input.PoolID,
strings.TrimRight(strings.TrimSpace(input.Endpoint), "/"),
labels,
capabilities,
strings.TrimSpace(input.ProtocolVersion),
strings.TrimSpace(input.OrchestratorInstanceRef),
strings.TrimSpace(input.Revision),
input.DesiredCapacity,
input.CapacityLimit,
@@ -235,6 +287,8 @@ WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
}
return WorkerAllocation{
InstanceID: input.InstanceID,
WorkerID: input.WorkerID,
PoolID: input.PoolID,
DesiredCapacity: input.DesiredCapacity,
Allocated: allocated,
GlobalAllocated: globalAllocated,
@@ -310,28 +364,30 @@ WHERE instance_id = $1
}
type WorkerInstanceRuntime struct {
InstanceID string `json:"instanceId"`
PodUID string `json:"podUid,omitempty"`
PodName string `json:"podName,omitempty"`
Site string `json:"site,omitempty"`
Revision string `json:"revision,omitempty"`
Status string `json:"status"`
Allocated int `json:"allocatedCapacity"`
CapacityLimit int `json:"capacityLimit"`
HardCapacityLimit int `json:"hardCapacityLimit"`
SafeCapacity int `json:"safeCapacity"`
HeavyCapacity int `json:"heavyCapacity"`
ReportedActiveTasks int `json:"reportedActiveTasks"`
PreparingTasks int `json:"preparingTasks"`
WaitingUpstreamTasks int `json:"waitingUpstreamTasks"`
FinalizingTasks int `json:"finalizingTasks"`
PressureState string `json:"pressureState"`
PressureReason string `json:"pressureReason,omitempty"`
LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"`
RunningTasks int `json:"runningTasks"`
ActiveLeases int `json:"activeLeases"`
HeartbeatAt time.Time `json:"heartbeatAt"`
DrainingAt *time.Time `json:"drainingAt,omitempty"`
InstanceID string `json:"instanceId"`
PodUID string `json:"podUid,omitempty"`
PodName string `json:"podName,omitempty"`
OrchestratorInstanceRef string `json:"-"`
Site string `json:"site,omitempty"`
PoolID string `json:"poolId,omitempty"`
Revision string `json:"revision,omitempty"`
Status string `json:"status"`
Allocated int `json:"allocatedCapacity"`
CapacityLimit int `json:"capacityLimit"`
HardCapacityLimit int `json:"hardCapacityLimit"`
SafeCapacity int `json:"safeCapacity"`
HeavyCapacity int `json:"heavyCapacity"`
ReportedActiveTasks int `json:"reportedActiveTasks"`
PreparingTasks int `json:"preparingTasks"`
WaitingUpstreamTasks int `json:"waitingUpstreamTasks"`
FinalizingTasks int `json:"finalizingTasks"`
PressureState string `json:"pressureState"`
PressureReason string `json:"pressureReason,omitempty"`
LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"`
RunningTasks int `json:"runningTasks"`
ActiveLeases int `json:"activeLeases"`
HeartbeatAt time.Time `json:"heartbeatAt"`
DrainingAt *time.Time `json:"drainingAt,omitempty"`
}
type WorkerQueueRuntime struct {
@@ -340,6 +396,12 @@ type WorkerQueueRuntime struct {
OldestWaitSeconds float64 `json:"oldestWaitSeconds"`
}
type PoolQueueRuntime struct {
PoolID string `json:"poolId"`
Queued int `json:"queued"`
Running int `json:"running"`
}
type WorkerClusterRuntime struct {
Workers []WorkerInstanceRuntime `json:"workers"`
Queue WorkerQueueRuntime `json:"queue"`
@@ -380,12 +442,40 @@ WHERE status IN ('queued', 'running')
return snapshot, err
}
func (s *Store) ListPoolQueueRuntime(ctx context.Context) ([]PoolQueueRuntime, error) {
rows, err := s.pool.Query(ctx, `
SELECT assigned_pool_id,
count(*) FILTER (WHERE status = 'queued')::int,
count(*) FILTER (WHERE status = 'running')::int
FROM gateway_tasks
WHERE assigned_pool_id IS NOT NULL
AND status IN ('queued', 'running')
AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
GROUP BY assigned_pool_id
ORDER BY assigned_pool_id`)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]PoolQueueRuntime, 0)
for rows.Next() {
var item PoolQueueRuntime
if err := rows.Scan(&item.PoolID, &item.Queued, &item.Running); err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) ListWorkerInstanceRuntime(ctx context.Context) ([]WorkerInstanceRuntime, error) {
rows, err := s.pool.Query(ctx, `
SELECT worker.instance_id,
worker.pod_uid,
worker.pod_name,
worker.orchestrator_instance_ref,
worker.site,
worker.pool_id,
worker.revision,
worker.status,
worker.allocated_capacity,
@@ -414,7 +504,7 @@ LEFT JOIN gateway_concurrency_leases lease
WHERE worker.status IN ('active', 'draining')
AND worker.heartbeat_at > now() - $1::interval
GROUP BY worker.instance_id
ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`,
ORDER BY worker.pool_id ASC, worker.status DESC, worker.instance_id ASC`,
runtimeWorkerStaleAfter.String(),
)
if err != nil {
@@ -428,7 +518,9 @@ ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`,
&instance.InstanceID,
&instance.PodUID,
&instance.PodName,
&instance.OrchestratorInstanceRef,
&instance.Site,
&instance.PoolID,
&instance.Revision,
&instance.Status,
&instance.Allocated,