wangbo
7786692d32
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
455 lines
15 KiB
Go
455 lines
15 KiB
Go
package capacitycontroller
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
const (
|
|
controllerReconcileInterval = 10 * time.Second
|
|
controllerLeadershipRetry = 5 * time.Second
|
|
controllerDeletionCost = -1000
|
|
)
|
|
|
|
type capacityStore interface {
|
|
TryAcquireCapacityControllerLeadership(context.Context) (store.Leadership, bool, error)
|
|
WorkerQueueRuntime(context.Context) (store.WorkerQueueRuntime, error)
|
|
ListPoolQueueRuntime(context.Context) ([]store.PoolQueueRuntime, error)
|
|
ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error)
|
|
CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error)
|
|
MarkWorkerDraining(context.Context, string) error
|
|
ReactivateWorkerInstance(context.Context, string) error
|
|
PublishDesiredCapacity(context.Context, executionpool.DesiredCapacity) error
|
|
}
|
|
|
|
type OrchestratorAdapter interface {
|
|
PoolState(context.Context, string, string) (PoolInfrastructureState, error)
|
|
ScalePool(context.Context, string, string, int) error
|
|
SetInstanceTerminationPriority(context.Context, string, int) error
|
|
}
|
|
|
|
type Status struct {
|
|
Leader bool `json:"leader"`
|
|
LastRunAt time.Time `json:"lastRunAt,omitempty"`
|
|
LastError string `json:"lastError,omitempty"`
|
|
Queue store.WorkerQueueRuntime `json:"queue"`
|
|
Plan Plan `json:"plan"`
|
|
ScaleActions uint64 `json:"scaleActions"`
|
|
}
|
|
|
|
type Controller struct {
|
|
cfg config.Config
|
|
store capacityStore
|
|
orchestrator OrchestratorAdapter
|
|
pools []config.ExecutionPoolCapacityConfig
|
|
logger *slog.Logger
|
|
expectedRevision string
|
|
now func() time.Time
|
|
highSince time.Time
|
|
lowSince time.Time
|
|
statusMu sync.RWMutex
|
|
status Status
|
|
}
|
|
|
|
func New(
|
|
cfg config.Config,
|
|
db capacityStore,
|
|
orchestrator OrchestratorAdapter,
|
|
logger *slog.Logger,
|
|
) *Controller {
|
|
return &Controller{
|
|
cfg: cfg, store: db, orchestrator: orchestrator, logger: logger,
|
|
expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
|
|
now: time.Now,
|
|
pools: mustCapacityPools(cfg),
|
|
}
|
|
}
|
|
|
|
func mustCapacityPools(cfg config.Config) []config.ExecutionPoolCapacityConfig {
|
|
pools, _ := cfg.CapacityPools()
|
|
return pools
|
|
}
|
|
|
|
func (controller *Controller) Run(ctx context.Context) {
|
|
for ctx.Err() == nil {
|
|
leadership, acquired, err := controller.store.TryAcquireCapacityControllerLeadership(ctx)
|
|
if err != nil {
|
|
controller.setError(err)
|
|
if !waitContext(ctx, controllerLeadershipRetry) {
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
if !acquired {
|
|
controller.setLeader(false)
|
|
controller.clearError()
|
|
if !waitContext(ctx, controllerLeadershipRetry) {
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
controller.setLeader(true)
|
|
controller.clearError()
|
|
controller.runLeader(ctx, leadership)
|
|
leadership.Release()
|
|
controller.setLeader(false)
|
|
}
|
|
}
|
|
|
|
func (controller *Controller) runLeader(ctx context.Context, leadership store.Leadership) {
|
|
ticker := time.NewTicker(controllerReconcileInterval)
|
|
defer ticker.Stop()
|
|
if err := controller.reconcile(ctx); err != nil {
|
|
controller.logError("initial capacity reconciliation failed", err)
|
|
}
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
keepAliveCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
|
err := leadership.KeepAlive(keepAliveCtx)
|
|
cancel()
|
|
if err != nil {
|
|
controller.logError("capacity controller leadership lost", err)
|
|
return
|
|
}
|
|
if err := controller.reconcile(ctx); err != nil {
|
|
controller.logError("capacity reconciliation failed", err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (controller *Controller) reconcile(ctx context.Context) error {
|
|
queue, err := controller.store.WorkerQueueRuntime(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
poolQueues, err := controller.store.ListPoolQueueRuntime(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
demandByPool := make(map[string]int, len(poolQueues))
|
|
for _, poolQueue := range poolQueues {
|
|
demandByPool[poolQueue.PoolID] = poolQueue.Queued + poolQueue.Running
|
|
}
|
|
instances, err := controller.store.ListWorkerInstanceRuntime(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
database, err := controller.store.CapacityDatabaseHealth(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now := controller.now()
|
|
poolResources := make([]PoolResources, 0, len(controller.pools))
|
|
infrastructureStates := make(map[string]PoolInfrastructureState, len(controller.pools))
|
|
for _, pool := range controller.pools {
|
|
if pool.MaxReplicas == 0 {
|
|
continue
|
|
}
|
|
state, stateErr := controller.orchestrator.PoolState(ctx, pool.ID, pool.AdapterRef)
|
|
if stateErr != nil {
|
|
return stateErr
|
|
}
|
|
infrastructureStates[pool.ID] = state
|
|
poolResources = append(poolResources, PoolResources{
|
|
PoolID: pool.ID, CurrentReplicas: state.CurrentReplicas,
|
|
MinReplicas: pool.MinReplicas, MaxReplicas: pool.MaxReplicas,
|
|
Demand: demandByPool[pool.ID],
|
|
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
|
|
UsedMemoryBytes: state.UsedMemoryBytes,
|
|
WorkerRequestMemoryBytes: state.WorkerRequestMemoryBytes,
|
|
AllocatableMilliCPU: state.AllocatableMilliCPU,
|
|
UsedMilliCPU: state.UsedMilliCPU,
|
|
WorkerRequestMilliCPU: state.WorkerRequestMilliCPU,
|
|
MemoryPressure: state.MemoryPressure,
|
|
Nodes: nodeResources(state.Nodes),
|
|
})
|
|
}
|
|
currentTotal := 0
|
|
for _, pool := range poolResources {
|
|
currentTotal += pool.CurrentReplicas
|
|
}
|
|
target := controller.cfg.WorkerTargetOutstandingPerReplica
|
|
if target < 1 {
|
|
target = 2 * controller.cfg.AsyncWorkerInstanceHardLimit
|
|
}
|
|
rawDesired := (max(queue.Queued+queue.Running, 0) + target - 1) / target
|
|
if rawDesired > currentTotal {
|
|
if controller.highSince.IsZero() {
|
|
controller.highSince = now
|
|
}
|
|
controller.lowSince = time.Time{}
|
|
} else if rawDesired < currentTotal && queue.Queued == 0 {
|
|
if controller.lowSince.IsZero() {
|
|
controller.lowSince = now
|
|
}
|
|
controller.highSince = time.Time{}
|
|
} else {
|
|
controller.highSince = time.Time{}
|
|
controller.lowSince = time.Time{}
|
|
}
|
|
revisionHealthy := controller.revisionsMatch(instances, infrastructureStates)
|
|
scaleUpEligible := !controller.highSince.IsZero() &&
|
|
now.Sub(controller.highSince) >= time.Duration(controller.cfg.WorkerScaleUpWindowSeconds)*time.Second &&
|
|
revisionHealthy
|
|
scaleDownEligible := !controller.lowSince.IsZero() &&
|
|
now.Sub(controller.lowSince) >= time.Duration(controller.cfg.WorkerScaleDownStabilizationSeconds)*time.Second &&
|
|
revisionHealthy
|
|
plan := CalculatePlan(PlanInput{
|
|
Queued: queue.Queued, Running: queue.Running,
|
|
InstanceSlots: controller.cfg.AsyncWorkerInstanceHardLimit,
|
|
TargetOutstandingPerReplica: target,
|
|
MemoryTargetPercent: controller.cfg.NodeMemoryTargetPercent,
|
|
MemoryHardPercent: controller.cfg.NodeMemoryHardPercent,
|
|
CPUTargetPercent: controller.cfg.NodeCPUTargetPercent,
|
|
DatabaseConnections: database.Connections,
|
|
DatabaseConnectionBudget: controller.cfg.PostgresConnectionBudget,
|
|
NonWorkerConnectionBudget: controller.cfg.PostgresNonWorkerConnectionBudget,
|
|
WorkerDatabasePoolMax: controller.cfg.WorkerDatabaseMaxConns,
|
|
SynchronousDatabasePeers: database.SynchronousPeers,
|
|
ScaleUpEligible: scaleUpEligible,
|
|
ScaleDownEligible: scaleDownEligible,
|
|
Pools: poolResources,
|
|
})
|
|
if !revisionHealthy && plan.FrozenReason == "" {
|
|
plan.FrozenReason = "release_revision_mismatch"
|
|
}
|
|
for _, poolPlan := range plan.Pools {
|
|
reason := "capacity_plan"
|
|
if plan.FrozenReason != "" {
|
|
reason = "frozen:" + plan.FrozenReason
|
|
}
|
|
if err := controller.store.PublishDesiredCapacity(ctx, executionpool.DesiredCapacity{
|
|
PoolID: poolPlan.PoolID, Desired: poolPlan.DesiredReplicas,
|
|
Reason: reason, ValidUntil: now.Add(2 * controllerReconcileInterval),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if controller.cfg.WorkerAutoscalingEnabled {
|
|
for _, poolPlan := range plan.Pools {
|
|
switch {
|
|
case poolPlan.DesiredReplicas > poolPlan.CurrentReplicas:
|
|
if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.DesiredReplicas); err != nil {
|
|
return err
|
|
}
|
|
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.DesiredReplicas, "scale_up")
|
|
case poolPlan.DesiredReplicas < poolPlan.CurrentReplicas:
|
|
if err := controller.reconcileScaleDown(ctx, now, poolPlan, instances); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
}
|
|
controller.statusMu.Lock()
|
|
controller.status.LastRunAt = now
|
|
controller.status.LastError = ""
|
|
controller.status.Queue = queue
|
|
controller.status.Plan = plan
|
|
controller.statusMu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func nodeResources(states []KubernetesNodeState) []NodeResources {
|
|
nodes := make([]NodeResources, 0, len(states))
|
|
for _, state := range states {
|
|
nodes = append(nodes, NodeResources{
|
|
NodeName: state.NodeName,
|
|
CurrentReplicas: state.CurrentReplicas,
|
|
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
|
|
UsedMemoryBytes: state.UsedMemoryBytes,
|
|
WorkerUsedMemoryBytes: state.WorkerUsedMemoryBytes,
|
|
AllocatableMilliCPU: state.AllocatableMilliCPU,
|
|
UsedMilliCPU: state.UsedMilliCPU,
|
|
WorkerUsedMilliCPU: state.WorkerUsedMilliCPU,
|
|
MemoryPressure: state.MemoryPressure,
|
|
})
|
|
}
|
|
return nodes
|
|
}
|
|
|
|
func (controller *Controller) reconcileScaleDown(
|
|
ctx context.Context,
|
|
now time.Time,
|
|
poolPlan PoolPlan,
|
|
instances []store.WorkerInstanceRuntime,
|
|
) error {
|
|
for _, instance := range instances {
|
|
if workerPoolID(instance) != poolPlan.PoolID || instance.Status != "draining" {
|
|
continue
|
|
}
|
|
if instance.RunningTasks == 0 && instance.ActiveLeases == 0 {
|
|
if strings.TrimSpace(instance.OrchestratorInstanceRef) == "" {
|
|
return errors.New("worker orchestrator instance reference is required for scale down")
|
|
}
|
|
if err := controller.orchestrator.SetInstanceTerminationPriority(ctx, instance.OrchestratorInstanceRef, controllerDeletionCost); err != nil {
|
|
return err
|
|
}
|
|
if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.CurrentReplicas-1); err != nil {
|
|
return err
|
|
}
|
|
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas-1, "drained_scale_down")
|
|
return nil
|
|
}
|
|
if instance.DrainingAt != nil &&
|
|
now.Sub(*instance.DrainingAt) >= time.Duration(controller.cfg.WorkerDrainTimeoutSeconds)*time.Second {
|
|
if err := controller.store.ReactivateWorkerInstance(ctx, instance.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
|
return err
|
|
}
|
|
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_timeout")
|
|
}
|
|
return nil
|
|
}
|
|
var candidate *store.WorkerInstanceRuntime
|
|
for index := range instances {
|
|
instance := &instances[index]
|
|
if workerPoolID(*instance) != poolPlan.PoolID || instance.Status != "active" {
|
|
continue
|
|
}
|
|
if candidate == nil ||
|
|
instance.RunningTasks+instance.ActiveLeases < candidate.RunningTasks+candidate.ActiveLeases {
|
|
candidate = instance
|
|
}
|
|
}
|
|
if candidate == nil {
|
|
return nil
|
|
}
|
|
if err := controller.store.MarkWorkerDraining(ctx, candidate.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
|
return err
|
|
}
|
|
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_started")
|
|
return nil
|
|
}
|
|
|
|
func workerPoolID(instance store.WorkerInstanceRuntime) string {
|
|
if strings.TrimSpace(instance.PoolID) != "" {
|
|
return instance.PoolID
|
|
}
|
|
return instance.Site
|
|
}
|
|
|
|
func (controller *Controller) adapterRef(poolID string) string {
|
|
for _, pool := range controller.pools {
|
|
if pool.ID == poolID {
|
|
return pool.AdapterRef
|
|
}
|
|
}
|
|
return poolID
|
|
}
|
|
|
|
func (controller *Controller) revisionsMatch(
|
|
instances []store.WorkerInstanceRuntime,
|
|
states map[string]PoolInfrastructureState,
|
|
) bool {
|
|
if controller.expectedRevision == "" {
|
|
return true
|
|
}
|
|
for _, state := range states {
|
|
if state.Revision != "" && state.Revision != controller.expectedRevision {
|
|
return false
|
|
}
|
|
}
|
|
for _, instance := range instances {
|
|
if instance.Revision != "" && instance.Revision != controller.expectedRevision {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (controller *Controller) Status() Status {
|
|
controller.statusMu.RLock()
|
|
defer controller.statusMu.RUnlock()
|
|
return controller.status
|
|
}
|
|
|
|
func (controller *Controller) Handler() http.Handler {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"ok":true,"service":"easyai-capacity-controller"}`))
|
|
})
|
|
mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, _ *http.Request) {
|
|
status := controller.Status()
|
|
// Followers are intentionally idle while the database advisory lock is
|
|
// held by the elected leader. They are still ready to take over and must
|
|
// not make a two-replica Deployment permanently fail its rollout.
|
|
if status.LastError != "" {
|
|
http.Error(w, `{"ok":false}`, http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"ok":true}`))
|
|
})
|
|
mux.HandleFunc("GET /status", func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(controller.Status())
|
|
})
|
|
return mux
|
|
}
|
|
|
|
func (controller *Controller) setLeader(leader bool) {
|
|
controller.statusMu.Lock()
|
|
controller.status.Leader = leader
|
|
controller.statusMu.Unlock()
|
|
}
|
|
|
|
func (controller *Controller) setError(err error) {
|
|
controller.statusMu.Lock()
|
|
controller.status.LastError = err.Error()
|
|
controller.statusMu.Unlock()
|
|
}
|
|
|
|
func (controller *Controller) clearError() {
|
|
controller.statusMu.Lock()
|
|
controller.status.LastError = ""
|
|
controller.statusMu.Unlock()
|
|
}
|
|
|
|
func (controller *Controller) logError(message string, err error) {
|
|
controller.setError(err)
|
|
if controller.logger != nil {
|
|
controller.logger.Error(message, "error", err)
|
|
}
|
|
}
|
|
|
|
func (controller *Controller) observeScale(site string, from int, to int, reason string) {
|
|
controller.statusMu.Lock()
|
|
controller.status.ScaleActions++
|
|
controller.statusMu.Unlock()
|
|
if controller.logger != nil {
|
|
controller.logger.Info("worker capacity action",
|
|
"site", site,
|
|
"fromReplicas", from,
|
|
"toReplicas", to,
|
|
"reason", reason,
|
|
)
|
|
}
|
|
}
|
|
|
|
func waitContext(ctx context.Context, duration time.Duration) bool {
|
|
timer := time.NewTimer(duration)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-ctx.Done():
|
|
return false
|
|
case <-timer.C:
|
|
return true
|
|
}
|
|
}
|