feat(acceptance): 建立同构验收与弹性容量体系
实现本地三节点 K3s 同构环境、脱敏生产快照、Gemini 图片和多参考图视频协议模拟、统一验收报告及故障注入。\n\n新增 Worker 容量控制器、资源与连接预算、任务恢复保护，并将生产验收拆分为 validation 执行和人工 CAS 放量。\n\n验证包括 Go 全量测试、PostgreSQL HTTP 集成测试、go vet、OpenAPI、ShellCheck、前端检查、迁移及发布脚本测试。
This commit is contained in:
@@ -0,0 +1,413 @@
|
||||
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/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)
|
||||
ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error)
|
||||
CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error)
|
||||
MarkWorkerDraining(context.Context, string) error
|
||||
ReactivateWorkerInstance(context.Context, string) error
|
||||
}
|
||||
|
||||
type capacityKubernetes interface {
|
||||
SiteState(context.Context, string) (KubernetesSiteState, error)
|
||||
ScaleWorkerDeployment(context.Context, string, int) error
|
||||
SetPodDeletionCost(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
|
||||
kubernetes capacityKubernetes
|
||||
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,
|
||||
kubernetes capacityKubernetes,
|
||||
logger *slog.Logger,
|
||||
) *Controller {
|
||||
return &Controller{
|
||||
cfg: cfg, store: db, kubernetes: kubernetes, logger: logger,
|
||||
expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
|
||||
now: time.Now,
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
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()
|
||||
sites := make([]SiteResources, 0, 2)
|
||||
kubernetesStates := make(map[string]KubernetesSiteState, 2)
|
||||
for _, site := range []string{"ningbo", "hongkong"} {
|
||||
state, stateErr := controller.kubernetes.SiteState(ctx, site)
|
||||
if stateErr != nil {
|
||||
return stateErr
|
||||
}
|
||||
kubernetesStates[site] = state
|
||||
minReplicas, maxReplicas := controller.siteReplicaBounds(site)
|
||||
sites = append(sites, SiteResources{
|
||||
Site: site, CurrentReplicas: state.CurrentReplicas,
|
||||
MinReplicas: minReplicas, MaxReplicas: maxReplicas,
|
||||
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 _, site := range sites {
|
||||
currentTotal += site.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, kubernetesStates)
|
||||
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,
|
||||
Sites: sites,
|
||||
})
|
||||
if !revisionHealthy && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "release_revision_mismatch"
|
||||
}
|
||||
if controller.cfg.WorkerAutoscalingEnabled {
|
||||
for _, sitePlan := range plan.Sites {
|
||||
switch {
|
||||
case sitePlan.DesiredReplicas > sitePlan.CurrentReplicas:
|
||||
if err := controller.kubernetes.ScaleWorkerDeployment(ctx, sitePlan.Site, sitePlan.DesiredReplicas); err != nil {
|
||||
return err
|
||||
}
|
||||
controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.DesiredReplicas, "scale_up")
|
||||
case sitePlan.DesiredReplicas < sitePlan.CurrentReplicas:
|
||||
if err := controller.reconcileScaleDown(ctx, now, sitePlan, 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,
|
||||
sitePlan SitePlan,
|
||||
instances []store.WorkerInstanceRuntime,
|
||||
) error {
|
||||
for _, instance := range instances {
|
||||
if instance.Site != sitePlan.Site || instance.Status != "draining" {
|
||||
continue
|
||||
}
|
||||
if instance.RunningTasks == 0 && instance.ActiveLeases == 0 {
|
||||
if err := controller.kubernetes.SetPodDeletionCost(ctx, instance.PodName, controllerDeletionCost); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := controller.kubernetes.ScaleWorkerDeployment(ctx, sitePlan.Site, sitePlan.CurrentReplicas-1); err != nil {
|
||||
return err
|
||||
}
|
||||
controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.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(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas, "drain_timeout")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
var candidate *store.WorkerInstanceRuntime
|
||||
for index := range instances {
|
||||
instance := &instances[index]
|
||||
if instance.Site != sitePlan.Site || 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(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas, "drain_started")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (controller *Controller) siteReplicaBounds(site string) (int, int) {
|
||||
switch site {
|
||||
case "ningbo":
|
||||
return controller.cfg.WorkerMinReplicasNingbo, controller.cfg.WorkerMaxReplicasNingbo
|
||||
case "hongkong":
|
||||
return controller.cfg.WorkerMinReplicasHongkong, controller.cfg.WorkerMaxReplicasHongkong
|
||||
default:
|
||||
return 0, 0
|
||||
}
|
||||
}
|
||||
|
||||
func (controller *Controller) revisionsMatch(
|
||||
instances []store.WorkerInstanceRuntime,
|
||||
states map[string]KubernetesSiteState,
|
||||
) 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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,440 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type KubernetesConfig struct {
|
||||
Namespace string
|
||||
APIServer string
|
||||
TokenFile string
|
||||
CAFile string
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
type KubernetesClient struct {
|
||||
namespace string
|
||||
baseURL string
|
||||
tokenFile string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
type KubernetesSiteState struct {
|
||||
Site string
|
||||
NodeName string
|
||||
CurrentReplicas int
|
||||
Revision string
|
||||
Image string
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerRequestMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerRequestMilliCPU int64
|
||||
MemoryPressure bool
|
||||
Nodes []KubernetesNodeState
|
||||
}
|
||||
|
||||
type KubernetesNodeState struct {
|
||||
NodeName string
|
||||
CurrentReplicas int
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerUsedMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerUsedMilliCPU int64
|
||||
MemoryPressure bool
|
||||
}
|
||||
|
||||
func NewKubernetesClient(config KubernetesConfig) (*KubernetesClient, error) {
|
||||
if strings.TrimSpace(config.Namespace) == "" {
|
||||
return nil, errors.New("capacity controller Kubernetes namespace is required")
|
||||
}
|
||||
if config.APIServer == "" {
|
||||
config.APIServer = "https://kubernetes.default.svc"
|
||||
}
|
||||
if config.TokenFile == "" {
|
||||
config.TokenFile = "/var/run/secrets/kubernetes.io/serviceaccount/token"
|
||||
}
|
||||
parsed, err := url.Parse(strings.TrimRight(config.APIServer, "/"))
|
||||
if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" ||
|
||||
(parsed.Scheme != "https" && config.HTTPClient == nil) {
|
||||
return nil, errors.New("capacity controller Kubernetes API server URL is invalid")
|
||||
}
|
||||
client := config.HTTPClient
|
||||
if client == nil {
|
||||
if config.CAFile == "" {
|
||||
config.CAFile = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
|
||||
}
|
||||
certificate, readErr := os.ReadFile(config.CAFile)
|
||||
if readErr != nil {
|
||||
return nil, fmt.Errorf("read Kubernetes service account CA: %w", readErr)
|
||||
}
|
||||
roots := x509.NewCertPool()
|
||||
if !roots.AppendCertsFromPEM(certificate) {
|
||||
return nil, errors.New("Kubernetes service account CA is invalid")
|
||||
}
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.Proxy = nil
|
||||
transport.DisableCompression = true
|
||||
transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
|
||||
client = &http.Client{
|
||||
Timeout: 10 * time.Second,
|
||||
Transport: transport,
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
}
|
||||
return &KubernetesClient{
|
||||
namespace: config.Namespace,
|
||||
baseURL: strings.TrimRight(config.APIServer, "/"),
|
||||
tokenFile: config.TokenFile,
|
||||
client: client,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) SiteState(ctx context.Context, site string) (KubernetesSiteState, error) {
|
||||
site = strings.TrimSpace(site)
|
||||
if site == "" {
|
||||
return KubernetesSiteState{}, errors.New("site is required")
|
||||
}
|
||||
var deployment struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"metadata"`
|
||||
Spec struct {
|
||||
Replicas int `json:"replicas"`
|
||||
Template struct {
|
||||
Spec struct {
|
||||
Containers []struct {
|
||||
Name string `json:"name"`
|
||||
Image string `json:"image"`
|
||||
Env []struct {
|
||||
Name string `json:"name"`
|
||||
Value string `json:"value"`
|
||||
} `json:"env"`
|
||||
Resources struct {
|
||||
Requests map[string]string `json:"requests"`
|
||||
} `json:"resources"`
|
||||
} `json:"containers"`
|
||||
} `json:"spec"`
|
||||
} `json:"template"`
|
||||
} `json:"spec"`
|
||||
}
|
||||
deploymentPath := fmt.Sprintf(
|
||||
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.PathEscape(site),
|
||||
)
|
||||
if err := client.getJSON(ctx, deploymentPath, &deployment); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
state := KubernetesSiteState{Site: site, CurrentReplicas: deployment.Spec.Replicas}
|
||||
for _, container := range deployment.Spec.Template.Spec.Containers {
|
||||
if container.Name != "worker" {
|
||||
continue
|
||||
}
|
||||
state.Image = container.Image
|
||||
state.WorkerRequestMemoryBytes, _ = parseBinaryQuantity(container.Resources.Requests["memory"])
|
||||
state.WorkerRequestMilliCPU, _ = parseCPUQuantity(container.Resources.Requests["cpu"])
|
||||
for _, environment := range container.Env {
|
||||
if environment.Name == "AI_GATEWAY_REVISION" {
|
||||
state.Revision = environment.Value
|
||||
}
|
||||
}
|
||||
}
|
||||
var nodes struct {
|
||||
Items []struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"metadata"`
|
||||
Status struct {
|
||||
Allocatable map[string]string `json:"allocatable"`
|
||||
Conditions []struct {
|
||||
Type string `json:"type"`
|
||||
Status string `json:"status"`
|
||||
} `json:"conditions"`
|
||||
} `json:"status"`
|
||||
} `json:"items"`
|
||||
}
|
||||
nodePath := "/api/v1/nodes?labelSelector=" + url.QueryEscape("easyai.io/site="+site)
|
||||
if err := client.getJSON(ctx, nodePath, &nodes); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
if len(nodes.Items) == 0 {
|
||||
return KubernetesSiteState{}, fmt.Errorf("site %s has no matching nodes", site)
|
||||
}
|
||||
var pods struct {
|
||||
Items []struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
DeletionTimestamp *time.Time `json:"deletionTimestamp"`
|
||||
} `json:"metadata"`
|
||||
Spec struct {
|
||||
NodeName string `json:"nodeName"`
|
||||
} `json:"spec"`
|
||||
} `json:"items"`
|
||||
}
|
||||
podPath := fmt.Sprintf(
|
||||
"/api/v1/namespaces/%s/pods?labelSelector=%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
|
||||
)
|
||||
if err := client.getJSON(ctx, podPath, &pods); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
replicasByNode := make(map[string]int, len(nodes.Items))
|
||||
nodeByPod := make(map[string]string, len(pods.Items))
|
||||
for _, pod := range pods.Items {
|
||||
if pod.Metadata.DeletionTimestamp == nil && pod.Spec.NodeName != "" {
|
||||
replicasByNode[pod.Spec.NodeName]++
|
||||
nodeByPod[pod.Metadata.Name] = pod.Spec.NodeName
|
||||
}
|
||||
}
|
||||
var podMetrics struct {
|
||||
Items []struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"metadata"`
|
||||
Containers []struct {
|
||||
Usage map[string]string `json:"usage"`
|
||||
} `json:"containers"`
|
||||
} `json:"items"`
|
||||
}
|
||||
podMetricsPath := fmt.Sprintf(
|
||||
"/apis/metrics.k8s.io/v1beta1/namespaces/%s/pods?labelSelector=%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
|
||||
)
|
||||
if err := client.getJSON(ctx, podMetricsPath, &podMetrics); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
workerMemoryByNode := make(map[string]int64, len(nodes.Items))
|
||||
workerCPUByNode := make(map[string]int64, len(nodes.Items))
|
||||
measuredPods := make(map[string]bool, len(podMetrics.Items))
|
||||
for _, pod := range podMetrics.Items {
|
||||
nodeName := nodeByPod[pod.Metadata.Name]
|
||||
if nodeName == "" {
|
||||
continue
|
||||
}
|
||||
measuredPods[pod.Metadata.Name] = true
|
||||
for _, container := range pod.Containers {
|
||||
memory, memoryErr := parseBinaryQuantity(container.Usage["memory"])
|
||||
cpu, cpuErr := parseCPUQuantity(container.Usage["cpu"])
|
||||
if memoryErr != nil || cpuErr != nil {
|
||||
return KubernetesSiteState{}, fmt.Errorf(
|
||||
"site %s pod %s returned incomplete resource metrics",
|
||||
site,
|
||||
pod.Metadata.Name,
|
||||
)
|
||||
}
|
||||
workerMemoryByNode[nodeName] += memory
|
||||
workerCPUByNode[nodeName] += cpu
|
||||
}
|
||||
}
|
||||
for podName := range nodeByPod {
|
||||
if !measuredPods[podName] {
|
||||
return KubernetesSiteState{}, fmt.Errorf(
|
||||
"site %s pod %s has no resource metrics yet",
|
||||
site,
|
||||
podName,
|
||||
)
|
||||
}
|
||||
}
|
||||
for _, node := range nodes.Items {
|
||||
nodeState := KubernetesNodeState{
|
||||
NodeName: node.Metadata.Name,
|
||||
CurrentReplicas: replicasByNode[node.Metadata.Name],
|
||||
WorkerUsedMemoryBytes: workerMemoryByNode[node.Metadata.Name],
|
||||
WorkerUsedMilliCPU: workerCPUByNode[node.Metadata.Name],
|
||||
}
|
||||
nodeState.AllocatableMemoryBytes, _ = parseBinaryQuantity(node.Status.Allocatable["memory"])
|
||||
nodeState.AllocatableMilliCPU, _ = parseCPUQuantity(node.Status.Allocatable["cpu"])
|
||||
for _, condition := range node.Status.Conditions {
|
||||
if condition.Type == "MemoryPressure" && condition.Status == "True" {
|
||||
nodeState.MemoryPressure = true
|
||||
state.MemoryPressure = true
|
||||
}
|
||||
}
|
||||
var metrics struct {
|
||||
Usage map[string]string `json:"usage"`
|
||||
}
|
||||
metricsPath := "/apis/metrics.k8s.io/v1beta1/nodes/" + url.PathEscape(nodeState.NodeName)
|
||||
if err := client.getJSON(ctx, metricsPath, &metrics); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
nodeState.UsedMemoryBytes, _ = parseBinaryQuantity(metrics.Usage["memory"])
|
||||
nodeState.UsedMilliCPU, _ = parseCPUQuantity(metrics.Usage["cpu"])
|
||||
if nodeState.AllocatableMemoryBytes <= 0 || nodeState.AllocatableMilliCPU <= 0 {
|
||||
return KubernetesSiteState{}, fmt.Errorf(
|
||||
"site %s node %s returned incomplete resource quantities",
|
||||
site,
|
||||
nodeState.NodeName,
|
||||
)
|
||||
}
|
||||
state.Nodes = append(state.Nodes, nodeState)
|
||||
state.AllocatableMemoryBytes += nodeState.AllocatableMemoryBytes
|
||||
state.AllocatableMilliCPU += nodeState.AllocatableMilliCPU
|
||||
state.UsedMemoryBytes += nodeState.UsedMemoryBytes
|
||||
state.UsedMilliCPU += nodeState.UsedMilliCPU
|
||||
}
|
||||
if len(state.Nodes) == 1 {
|
||||
state.NodeName = state.Nodes[0].NodeName
|
||||
}
|
||||
if state.WorkerRequestMemoryBytes <= 0 || state.WorkerRequestMilliCPU <= 0 ||
|
||||
state.AllocatableMemoryBytes <= 0 || state.AllocatableMilliCPU <= 0 {
|
||||
return KubernetesSiteState{}, fmt.Errorf("site %s returned incomplete resource quantities", site)
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) ScaleWorkerDeployment(ctx context.Context, site string, replicas int) error {
|
||||
if replicas < 0 || replicas > 64 {
|
||||
return errors.New("worker replica target must be between 0 and 64")
|
||||
}
|
||||
path := fmt.Sprintf(
|
||||
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s/scale",
|
||||
url.PathEscape(client.namespace),
|
||||
url.PathEscape(strings.TrimSpace(site)),
|
||||
)
|
||||
return client.patchJSON(ctx, path, map[string]any{"spec": map[string]any{"replicas": replicas}})
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) SetPodDeletionCost(ctx context.Context, podName string, cost int) error {
|
||||
podName = strings.TrimSpace(podName)
|
||||
if podName == "" {
|
||||
return errors.New("worker pod name is required")
|
||||
}
|
||||
path := fmt.Sprintf(
|
||||
"/api/v1/namespaces/%s/pods/%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.PathEscape(podName),
|
||||
)
|
||||
return client.patchJSON(ctx, path, map[string]any{
|
||||
"metadata": map[string]any{
|
||||
"annotations": map[string]any{
|
||||
"controller.kubernetes.io/pod-deletion-cost": strconv.Itoa(cost),
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) getJSON(ctx context.Context, path string, target any) error {
|
||||
request, err := client.request(ctx, http.MethodGet, path, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
response, err := client.client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("request Kubernetes API: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
|
||||
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
|
||||
}
|
||||
decoder := json.NewDecoder(io.LimitReader(response.Body, 4*1024*1024))
|
||||
if err := decoder.Decode(target); err != nil {
|
||||
return fmt.Errorf("decode Kubernetes API %s: %w", path, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) patchJSON(ctx context.Context, path string, payload any) error {
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request, err := client.request(ctx, http.MethodPatch, path, bytes.NewReader(encoded))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/merge-patch+json")
|
||||
response, err := client.client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("patch Kubernetes API: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) request(ctx context.Context, method string, path string, body io.Reader) (*http.Request, error) {
|
||||
token, err := os.ReadFile(client.tokenFile)
|
||||
if err != nil || strings.TrimSpace(string(token)) == "" {
|
||||
clear(token)
|
||||
return nil, errors.New("Kubernetes service account token is unavailable")
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, method, client.baseURL+path, body)
|
||||
if err != nil {
|
||||
clear(token)
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(token)))
|
||||
request.Header.Set("Accept", "application/json")
|
||||
clear(token)
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func parseBinaryQuantity(value string) (int64, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return 0, errors.New("quantity is empty")
|
||||
}
|
||||
units := []struct {
|
||||
suffix string
|
||||
multiplier int64
|
||||
}{
|
||||
{"Ei", 1 << 60},
|
||||
{"Pi", 1 << 50},
|
||||
{"Ti", 1 << 40},
|
||||
{"Gi", 1 << 30},
|
||||
{"Mi", 1 << 20},
|
||||
{"Ki", 1 << 10},
|
||||
{"G", 1_000_000_000},
|
||||
{"M", 1_000_000},
|
||||
{"K", 1_000},
|
||||
}
|
||||
for _, unit := range units {
|
||||
if strings.HasSuffix(value, unit.suffix) {
|
||||
number, err := strconv.ParseFloat(strings.TrimSuffix(value, unit.suffix), 64)
|
||||
return int64(number * float64(unit.multiplier)), err
|
||||
}
|
||||
}
|
||||
return strconv.ParseInt(value, 10, 64)
|
||||
}
|
||||
|
||||
func parseCPUQuantity(value string) (int64, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if strings.HasSuffix(value, "n") {
|
||||
nanocores, err := strconv.ParseInt(strings.TrimSuffix(value, "n"), 10, 64)
|
||||
return nanocores / 1_000_000, err
|
||||
}
|
||||
if strings.HasSuffix(value, "u") {
|
||||
microcores, err := strconv.ParseInt(strings.TrimSuffix(value, "u"), 10, 64)
|
||||
return microcores / 1_000, err
|
||||
}
|
||||
if strings.HasSuffix(value, "m") {
|
||||
return strconv.ParseInt(strings.TrimSuffix(value, "m"), 10, 64)
|
||||
}
|
||||
cores, err := strconv.ParseFloat(value, 64)
|
||||
return int64(cores * 1000), err
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestKubernetesClientReadsSiteAndMutatesOnlyWorkerScale(t *testing.T) {
|
||||
var patches []string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.Header.Get("Authorization") != "Bearer test-token" {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-hongkong"):
|
||||
_, _ = w.Write([]byte(`{
|
||||
"spec":{"replicas":1,"template":{"spec":{"containers":[{
|
||||
"name":"worker","image":"registry.invalid/gateway@sha256:abc",
|
||||
"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
|
||||
"resources":{"requests":{"memory":"512Mi","cpu":"250m"}}
|
||||
}]}}}
|
||||
}`))
|
||||
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
|
||||
_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"easyai-hongkong"},"status":{
|
||||
"allocatable":{"memory":"8Gi","cpu":"4"},
|
||||
"conditions":[{"type":"MemoryPressure","status":"False"}]
|
||||
}}]}`))
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/api/v1/namespaces/easyai/pods"):
|
||||
_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"worker-hongkong-1"},"spec":{"nodeName":"easyai-hongkong"}}]}`))
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
|
||||
strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong"):
|
||||
_, _ = w.Write([]byte(`{"usage":{"memory":"3Gi","cpu":"500m"}}`))
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
|
||||
strings.HasSuffix(request.URL.Path, "/pods"):
|
||||
_, _ = w.Write([]byte(`{"items":[{
|
||||
"metadata":{"name":"worker-hongkong-1"},
|
||||
"containers":[{"usage":{"memory":"384Mi","cpu":"125m"}}]
|
||||
}]}`))
|
||||
case request.Method == http.MethodPatch:
|
||||
patches = append(patches, request.URL.Path)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
default:
|
||||
http.NotFound(w, request)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
tokenFile := filepath.Join(t.TempDir(), "token")
|
||||
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
client, err := NewKubernetesClient(KubernetesConfig{
|
||||
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
state, err := client.SiteState(context.Background(), "hongkong")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state.CurrentReplicas != 1 || state.NodeName != "easyai-hongkong" ||
|
||||
state.WorkerRequestMemoryBytes != 512<<20 || state.WorkerRequestMilliCPU != 250 ||
|
||||
state.AllocatableMemoryBytes != 8<<30 || state.UsedMemoryBytes != 3<<30 ||
|
||||
state.Nodes[0].WorkerUsedMemoryBytes != 384<<20 || state.Nodes[0].WorkerUsedMilliCPU != 125 {
|
||||
t.Fatalf("site state=%+v", state)
|
||||
}
|
||||
if err := client.SetPodDeletionCost(context.Background(), "worker-pod", -1000); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := client.ScaleWorkerDeployment(context.Background(), "hongkong", 2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(patches) != 2 ||
|
||||
!strings.Contains(patches[0], "/pods/worker-pod") ||
|
||||
!strings.Contains(patches[1], "/deployments/easyai-worker-hongkong/scale") {
|
||||
t.Fatalf("patches=%v", patches)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuantityParsing(t *testing.T) {
|
||||
if value, err := parseBinaryQuantity("1536Mi"); err != nil || value != 1536<<20 {
|
||||
t.Fatalf("memory=%d err=%v", value, err)
|
||||
}
|
||||
if value, err := parseCPUQuantity("250000000n"); err != nil || value != 250 {
|
||||
t.Fatalf("cpu=%d err=%v", value, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
)
|
||||
|
||||
type SiteResources struct {
|
||||
Site string
|
||||
CurrentReplicas int
|
||||
MinReplicas int
|
||||
MaxReplicas int
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerRequestMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerRequestMilliCPU int64
|
||||
MemoryPressure bool
|
||||
Nodes []NodeResources
|
||||
}
|
||||
|
||||
type NodeResources struct {
|
||||
NodeName string
|
||||
CurrentReplicas int
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerUsedMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerUsedMilliCPU int64
|
||||
MemoryPressure bool
|
||||
}
|
||||
|
||||
type PlanInput struct {
|
||||
Queued int
|
||||
Running int
|
||||
InstanceSlots int
|
||||
TargetOutstandingPerReplica int
|
||||
MemoryTargetPercent int
|
||||
MemoryHardPercent int
|
||||
CPUTargetPercent int
|
||||
DatabaseConnections int
|
||||
DatabaseConnectionBudget int
|
||||
NonWorkerConnectionBudget int
|
||||
WorkerDatabasePoolMax int
|
||||
SynchronousDatabasePeers int
|
||||
ScaleUpEligible bool
|
||||
ScaleDownEligible bool
|
||||
Sites []SiteResources
|
||||
}
|
||||
|
||||
type SitePlan struct {
|
||||
Site string `json:"site"`
|
||||
CurrentReplicas int `json:"currentReplicas"`
|
||||
DesiredReplicas int `json:"desiredReplicas"`
|
||||
ResourceMax int `json:"resourceMax"`
|
||||
MemoryPercent float64 `json:"memoryPercent"`
|
||||
CPUPercent float64 `json:"cpuPercent"`
|
||||
}
|
||||
|
||||
type Plan struct {
|
||||
RawDesired int `json:"rawDesired"`
|
||||
DesiredTotal int `json:"desiredTotal"`
|
||||
CurrentTotal int `json:"currentTotal"`
|
||||
FrozenReason string `json:"frozenReason,omitempty"`
|
||||
Sites []SitePlan `json:"sites"`
|
||||
}
|
||||
|
||||
func CalculatePlan(input PlanInput) Plan {
|
||||
if input.InstanceSlots < 1 {
|
||||
input.InstanceSlots = 1
|
||||
}
|
||||
target := input.TargetOutstandingPerReplica
|
||||
if target < 1 {
|
||||
target = 2 * input.InstanceSlots
|
||||
}
|
||||
if input.MemoryTargetPercent < 1 {
|
||||
input.MemoryTargetPercent = 75
|
||||
}
|
||||
if input.MemoryHardPercent < 1 {
|
||||
input.MemoryHardPercent = 85
|
||||
}
|
||||
if input.CPUTargetPercent < 1 {
|
||||
input.CPUTargetPercent = 70
|
||||
}
|
||||
rawDesired := int(math.Ceil(float64(max(input.Queued+input.Running, 0)) / float64(target)))
|
||||
plan := Plan{RawDesired: rawDesired}
|
||||
minTotal := 0
|
||||
resourceMaxTotal := 0
|
||||
for _, site := range input.Sites {
|
||||
for _, node := range site.Nodes {
|
||||
if node.MemoryPressure {
|
||||
site.MemoryPressure = true
|
||||
}
|
||||
}
|
||||
current := max(site.CurrentReplicas, 0)
|
||||
minReplicas := max(site.MinReplicas, 0)
|
||||
configMax := max(site.MaxReplicas, minReplicas)
|
||||
resourceMax, memoryPercent, cpuPercent := siteResourceMaximum(
|
||||
site,
|
||||
input.MemoryTargetPercent,
|
||||
input.CPUTargetPercent,
|
||||
)
|
||||
resourceMax = min(resourceMax, configMax)
|
||||
if resourceMax < minReplicas && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "site_resource_budget"
|
||||
}
|
||||
resourceMax = max(resourceMax, minReplicas)
|
||||
plan.Sites = append(plan.Sites, SitePlan{
|
||||
Site: site.Site, CurrentReplicas: current, DesiredReplicas: minReplicas,
|
||||
ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent,
|
||||
})
|
||||
plan.CurrentTotal += current
|
||||
minTotal += minReplicas
|
||||
resourceMaxTotal += resourceMax
|
||||
if site.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) {
|
||||
plan.FrozenReason = "node_memory_pressure"
|
||||
}
|
||||
}
|
||||
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
|
||||
workerReplicaBudget := max(
|
||||
(input.DatabaseConnectionBudget-input.NonWorkerConnectionBudget)/input.WorkerDatabasePoolMax,
|
||||
0,
|
||||
)
|
||||
capSiteResourceMaxima(plan.Sites, workerReplicaBudget)
|
||||
resourceMaxTotal = 0
|
||||
for _, site := range plan.Sites {
|
||||
resourceMaxTotal += site.ResourceMax
|
||||
}
|
||||
if workerReplicaBudget < minTotal && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "database_connection_budget"
|
||||
}
|
||||
}
|
||||
desired := max(rawDesired, minTotal)
|
||||
desired = min(desired, resourceMaxTotal)
|
||||
|
||||
if input.SynchronousDatabasePeers < 1 && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "database_not_synchronous"
|
||||
}
|
||||
if desired > plan.CurrentTotal {
|
||||
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
|
||||
additional := max(
|
||||
(input.DatabaseConnectionBudget-input.DatabaseConnections)/input.WorkerDatabasePoolMax,
|
||||
0,
|
||||
)
|
||||
desired = min(desired, plan.CurrentTotal+additional)
|
||||
}
|
||||
if !input.ScaleUpEligible || plan.FrozenReason != "" ||
|
||||
(input.DatabaseConnectionBudget > 0 && input.DatabaseConnections >= input.DatabaseConnectionBudget) {
|
||||
desired = plan.CurrentTotal
|
||||
if plan.FrozenReason == "" {
|
||||
if !input.ScaleUpEligible {
|
||||
plan.FrozenReason = "scale_up_window"
|
||||
} else {
|
||||
plan.FrozenReason = "database_connection_budget"
|
||||
}
|
||||
}
|
||||
} else {
|
||||
desired = min(desired, max(plan.CurrentTotal*2, plan.CurrentTotal+2))
|
||||
}
|
||||
} else if desired < plan.CurrentTotal && !input.ScaleDownEligible {
|
||||
desired = plan.CurrentTotal
|
||||
if plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "scale_down_stabilization"
|
||||
}
|
||||
}
|
||||
desired = max(desired, minTotal)
|
||||
plan.DesiredTotal = desired
|
||||
distributeDesiredReplicas(plan.Sites, desired)
|
||||
return plan
|
||||
}
|
||||
|
||||
func capSiteResourceMaxima(sites []SitePlan, total int) {
|
||||
minimum := 0
|
||||
original := make(map[string]int, len(sites))
|
||||
for index := range sites {
|
||||
original[sites[index].Site] = sites[index].ResourceMax
|
||||
sites[index].ResourceMax = sites[index].DesiredReplicas
|
||||
minimum += sites[index].ResourceMax
|
||||
}
|
||||
target := max(total, minimum)
|
||||
assigned := minimum
|
||||
for assigned < target {
|
||||
sort.SliceStable(sites, func(left, right int) bool {
|
||||
leftRoom := original[sites[left].Site] - sites[left].ResourceMax
|
||||
rightRoom := original[sites[right].Site] - sites[right].ResourceMax
|
||||
if leftRoom != rightRoom {
|
||||
return leftRoom > rightRoom
|
||||
}
|
||||
return sites[left].Site < sites[right].Site
|
||||
})
|
||||
if original[sites[0].Site] <= sites[0].ResourceMax {
|
||||
break
|
||||
}
|
||||
sites[0].ResourceMax++
|
||||
assigned++
|
||||
}
|
||||
}
|
||||
|
||||
func siteResourceMaximum(site SiteResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) {
|
||||
if len(site.Nodes) > 0 {
|
||||
if site.WorkerRequestMemoryBytes <= 0 || site.WorkerRequestMilliCPU <= 0 {
|
||||
return 0, 0, 0
|
||||
}
|
||||
memoryMax := 0
|
||||
cpuMax := 0
|
||||
memoryPercent := float64(0)
|
||||
cpuPercent := float64(0)
|
||||
for _, node := range site.Nodes {
|
||||
nodeMemoryPercent := usagePercent(node.UsedMemoryBytes, node.AllocatableMemoryBytes)
|
||||
nodeCPUPercent := usagePercent(node.UsedMilliCPU, node.AllocatableMilliCPU)
|
||||
memoryPercent = max(memoryPercent, nodeMemoryPercent)
|
||||
cpuPercent = max(cpuPercent, nodeCPUPercent)
|
||||
nonWorkerMemory := max(
|
||||
node.UsedMemoryBytes-node.WorkerUsedMemoryBytes,
|
||||
0,
|
||||
)
|
||||
memoryBudget := node.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorkerMemory
|
||||
memoryMax += int(max(memoryBudget, 0) / site.WorkerRequestMemoryBytes)
|
||||
nonWorkerCPU := max(
|
||||
node.UsedMilliCPU-node.WorkerUsedMilliCPU,
|
||||
0,
|
||||
)
|
||||
cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU
|
||||
cpuMax += int(max(cpuBudget, 0) / site.WorkerRequestMilliCPU)
|
||||
}
|
||||
return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
|
||||
}
|
||||
memoryPercent := usagePercent(site.UsedMemoryBytes, site.AllocatableMemoryBytes)
|
||||
cpuPercent := usagePercent(site.UsedMilliCPU, site.AllocatableMilliCPU)
|
||||
memoryMax := site.MaxReplicas
|
||||
if site.AllocatableMemoryBytes > 0 && site.WorkerRequestMemoryBytes > 0 {
|
||||
nonWorker := max(site.UsedMemoryBytes-int64(site.CurrentReplicas)*site.WorkerRequestMemoryBytes, 0)
|
||||
budget := site.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker
|
||||
memoryMax = int(max(budget, 0) / site.WorkerRequestMemoryBytes)
|
||||
}
|
||||
cpuMax := site.MaxReplicas
|
||||
if site.AllocatableMilliCPU > 0 && site.WorkerRequestMilliCPU > 0 {
|
||||
nonWorker := max(site.UsedMilliCPU-int64(site.CurrentReplicas)*site.WorkerRequestMilliCPU, 0)
|
||||
budget := site.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker
|
||||
cpuMax = int(max(budget, 0) / site.WorkerRequestMilliCPU)
|
||||
}
|
||||
return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
|
||||
}
|
||||
|
||||
func distributeDesiredReplicas(sites []SitePlan, desired int) {
|
||||
if len(sites) == 0 {
|
||||
return
|
||||
}
|
||||
assigned := 0
|
||||
for index := range sites {
|
||||
assigned += sites[index].DesiredReplicas
|
||||
}
|
||||
for assigned < desired {
|
||||
sort.SliceStable(sites, func(left, right int) bool {
|
||||
leftRoom := sites[left].ResourceMax - sites[left].DesiredReplicas
|
||||
rightRoom := sites[right].ResourceMax - sites[right].DesiredReplicas
|
||||
if leftRoom != rightRoom {
|
||||
return leftRoom > rightRoom
|
||||
}
|
||||
if sites[left].DesiredReplicas != sites[right].DesiredReplicas {
|
||||
return sites[left].DesiredReplicas < sites[right].DesiredReplicas
|
||||
}
|
||||
return sites[left].Site < sites[right].Site
|
||||
})
|
||||
if sites[0].DesiredReplicas >= sites[0].ResourceMax {
|
||||
break
|
||||
}
|
||||
sites[0].DesiredReplicas++
|
||||
assigned++
|
||||
}
|
||||
sort.Slice(sites, func(left, right int) bool { return sites[left].Site < sites[right].Site })
|
||||
}
|
||||
|
||||
func usagePercent(used int64, allocatable int64) float64 {
|
||||
if allocatable <= 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(used) * 100 / float64(allocatable)
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) {
|
||||
input := PlanInput{
|
||||
Queued: 1000, Running: 48, InstanceSlots: 24,
|
||||
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
|
||||
DatabaseConnections: 80, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Sites: []SiteResources{
|
||||
{
|
||||
Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 6 << 30,
|
||||
WorkerRequestMemoryBytes: 1 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 2000, WorkerRequestMilliCPU: 500,
|
||||
},
|
||||
{
|
||||
Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
|
||||
WorkerRequestMemoryBytes: 1 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 500, WorkerRequestMilliCPU: 500,
|
||||
},
|
||||
},
|
||||
}
|
||||
plan := CalculatePlan(input)
|
||||
if plan.RawDesired != 22 {
|
||||
t.Fatalf("raw desired=%d, want 22", plan.RawDesired)
|
||||
}
|
||||
if plan.DesiredTotal != 4 {
|
||||
t.Fatalf("desired total=%d, want step-limited 4: %+v", plan.DesiredTotal, plan)
|
||||
}
|
||||
if plan.Sites[0].Site != "hongkong" || plan.Sites[0].DesiredReplicas != 3 {
|
||||
t.Fatalf("site allocation=%+v, want hongkong=3 ningbo=1", plan.Sites)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanFreezesScaleUpOnHardMemoryPressure(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 200, InstanceSlots: 24,
|
||||
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Sites: []SiteResources{{
|
||||
Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 7 << 30,
|
||||
WorkerRequestMemoryBytes: 1 << 30,
|
||||
}},
|
||||
})
|
||||
if plan.DesiredTotal != 1 || plan.FrozenReason != "node_memory_pressure" {
|
||||
t.Fatalf("plan=%+v, want frozen at one replica", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanReportsConfiguredMinimumOutsideResourceBudget(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 200, InstanceSlots: 24,
|
||||
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Sites: []SiteResources{{
|
||||
Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 13 << 29,
|
||||
WorkerRequestMemoryBytes: 2 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 2800, WorkerRequestMilliCPU: 500,
|
||||
}},
|
||||
})
|
||||
if plan.DesiredTotal != 1 || plan.FrozenReason != "site_resource_budget" {
|
||||
t.Fatalf("plan=%+v, want the existing minimum preserved and expansion frozen", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanWaitsForSafeScaleDownWindow(t *testing.T) {
|
||||
input := PlanInput{
|
||||
InstanceSlots: 24, DatabaseConnections: 20, DatabaseConnectionBudget: 150,
|
||||
WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1,
|
||||
ScaleDownEligible: false,
|
||||
Sites: []SiteResources{
|
||||
{Site: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
|
||||
{Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
|
||||
},
|
||||
}
|
||||
plan := CalculatePlan(input)
|
||||
if plan.DesiredTotal != 4 || plan.FrozenReason != "scale_down_stabilization" {
|
||||
t.Fatalf("plan=%+v, want current replicas during stabilization", plan)
|
||||
}
|
||||
input.ScaleDownEligible = true
|
||||
plan = CalculatePlan(input)
|
||||
if plan.DesiredTotal != 2 {
|
||||
t.Fatalf("desired total=%d, want min total 2", plan.DesiredTotal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 500, InstanceSlots: 24,
|
||||
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
|
||||
DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 20,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Sites: []SiteResources{{
|
||||
Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8,
|
||||
WorkerRequestMemoryBytes: 1 << 30, WorkerRequestMilliCPU: 500,
|
||||
Nodes: []NodeResources{
|
||||
{
|
||||
NodeName: "hk-a", CurrentReplicas: 1,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 1000,
|
||||
},
|
||||
{
|
||||
NodeName: "hk-b", CurrentReplicas: 1,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 2 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 500,
|
||||
},
|
||||
},
|
||||
}},
|
||||
})
|
||||
if plan.Sites[0].ResourceMax != 7 {
|
||||
t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Sites[0].ResourceMax, plan)
|
||||
}
|
||||
if plan.DesiredTotal != 4 {
|
||||
t.Fatalf("desired=%d, want two-wave step from 2 to 4", plan.DesiredTotal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 500, InstanceSlots: 24,
|
||||
DatabaseConnections: 20, DatabaseConnectionBudget: 150,
|
||||
NonWorkerConnectionBudget: 72, WorkerDatabasePoolMax: 32,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Sites: []SiteResources{
|
||||
{Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
|
||||
{Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
|
||||
},
|
||||
})
|
||||
if plan.DesiredTotal != 2 {
|
||||
t.Fatalf("desired=%d, want the current 1+1 database-safe topology: %+v", plan.DesiredTotal, plan)
|
||||
}
|
||||
resourceMax := 0
|
||||
for _, site := range plan.Sites {
|
||||
resourceMax += site.ResourceMax
|
||||
}
|
||||
if resourceMax != 2 {
|
||||
t.Fatalf("resource maximum=%d, want floor((150-72)/32)=2: %+v", resourceMax, plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanJSONUsesAcceptanceContractFieldNames(t *testing.T) {
|
||||
payload, err := json.Marshal(Plan{
|
||||
RawDesired: 3,
|
||||
Sites: []SitePlan{{Site: "hongkong", ResourceMax: 2}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
text := string(payload)
|
||||
for _, field := range []string{`"rawDesired":3`, `"sites"`, `"site":"hongkong"`, `"resourceMax":2`} {
|
||||
if !strings.Contains(text, field) {
|
||||
t.Fatalf("plan JSON %s does not contain %s", text, field)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user