perf(worker): 按实例内存上限扩展异步容量

将策略推导出的集群目标与单 Worker 内存安全上限分离,避免失活实例的全部容量转移到单个 2 GiB Pod。

生产环境单实例上限和媒体物化并发设为 24;两实例正常时总容量由 16 提升至 48,后续可通过增加 Worker 节点继续扩展。新增 0093 迁移保存实例容量上限,并兼容滚动发布中的旧实例。

风险:更高媒体并发会增加 Worker 和节点内存压力,保留单实例 24 的硬边界并由持续监控观察 RSS、OOM、队列与节点余量。

验证:go test ./... -count=1;go vet ./...;真实 PostgreSQL 分配与故障转移测试;迁移安全检查;kubectl kustomize;gofmt -l;git diff --check。
This commit is contained in:
2026-07-30 15:08:54 +08:00
parent ea58d21d03
commit 3aa78d3e5c
12 changed files with 180 additions and 25 deletions
+1 -1
View File
@@ -162,7 +162,7 @@ AI_GATEWAY_DATABASE_URL=postgresql://easyai:easyai2025@localhost:5432/easyai_ai_
如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。
异步队列 worker 不使用固定并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules``AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`仅作为单进程资源安全边界;平台模型和用户组的 PostgreSQL concurrency lease 是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
异步队列 worker 不使用固定业务并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的集群容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules``AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`限制策略推导出的集群目标;`AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT` 默认 `32`,按单 Worker 的内存安全容量限制实例分配,即使其他实例失活也不会突破。平台模型和用户组的 PostgreSQL concurrency lease 是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
## 迁移原则
+5
View File
@@ -63,6 +63,7 @@ type Config struct {
MediaMaterializationConcurrency int
AsyncQueueWorkerEnabled bool
AsyncWorkerHardLimit int
AsyncWorkerInstanceHardLimit int
AsyncWorkerRefreshIntervalSeconds int
}
@@ -123,6 +124,7 @@ func Load() Config {
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
AsyncWorkerInstanceHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", 32),
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
}
}
@@ -150,6 +152,9 @@ func (c Config) Validate() error {
if c.AsyncWorkerHardLimit < 1 || c.AsyncWorkerHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT must be between 1 and 10000")
}
if c.AsyncWorkerInstanceHardLimit < 1 || c.AsyncWorkerInstanceHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT must be between 1 and 10000")
}
if c.AsyncWorkerRefreshIntervalSeconds < 1 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
}
+24 -2
View File
@@ -26,7 +26,12 @@ func TestLoadIdentitySecretStoreUsesNewEnvironmentNamesAndIgnoresLegacyBusinessV
}
func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
cfg := Config{IdentitySecretStore: "file", AsyncWorkerHardLimit: 2048, AsyncWorkerRefreshIntervalSeconds: 5}
cfg := Config{
IdentitySecretStore: "file",
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") {
t.Fatalf("Validate() error = %v, want missing identity secret directory", err)
}
@@ -41,6 +46,7 @@ func TestValidateIdentityKubernetesSecretStore(t *testing.T) {
IdentitySecretStore: "kubernetes",
IdentityKubernetesSecretName: "easyai-gateway-identity",
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") {
@@ -58,6 +64,7 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
IdentitySecurityEventStaleAfterSeconds: 60,
IdentitySecurityEventClockSkewSeconds: 60,
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") {
@@ -66,11 +73,20 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
}
func TestValidateAsyncWorkerSettings(t *testing.T) {
cfg := Config{AsyncWorkerHardLimit: 10001, AsyncWorkerRefreshIntervalSeconds: 5}
cfg := Config{
AsyncWorkerHardLimit: 10001,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid hard limit", err)
}
cfg.AsyncWorkerHardLimit = 2048
cfg.AsyncWorkerInstanceHardLimit = 10001
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid instance hard limit", err)
}
cfg.AsyncWorkerInstanceHardLimit = 32
cfg.AsyncWorkerRefreshIntervalSeconds = 0
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid refresh interval", err)
@@ -80,6 +96,12 @@ func TestValidateAsyncWorkerSettings(t *testing.T) {
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
}
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "2048")
t.Setenv("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", "not-an-integer")
loaded = Load()
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer instance hard limit", err)
}
}
func TestLoadAsyncQueueWorkerEnabled(t *testing.T) {
+8 -1
View File
@@ -220,6 +220,10 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
s.logger.Info("async worker capacity initialized",
"capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired,
"activeInstances", snapshot.ActiveInstances,
"globalCapacity", snapshot.GlobalCapacity,
"instanceID", snapshot.InstanceID,
"instanceHardLimit", s.cfg.AsyncWorkerInstanceHardLimit,
"hardLimit", snapshot.HardLimit,
"enabledModels", snapshot.EnabledModels,
"unlimitedModels", snapshot.UnlimitedModels,
@@ -289,13 +293,14 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")),
Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
DesiredCapacity: snapshot.Capacity,
CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit,
HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second,
})
if err != nil {
return store.AsyncWorkerCapacitySnapshot{}, err
}
snapshot.Capacity = allocation.Allocated
snapshot.GlobalCapacity = allocation.DesiredCapacity
snapshot.GlobalCapacity = allocation.GlobalAllocated
snapshot.ActiveInstances = allocation.ActiveInstances
snapshot.InstanceID = allocation.InstanceID
return snapshot, nil
@@ -372,7 +377,9 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
"capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired,
"activeInstances", snapshot.ActiveInstances,
"globalCapacity", snapshot.GlobalCapacity,
"instanceID", snapshot.InstanceID,
"instanceHardLimit", s.cfg.AsyncWorkerInstanceHardLimit,
"hardLimit", snapshot.HardLimit,
"modelDesired", snapshot.ModelDesired,
"groupDesired", snapshot.GroupDesired,
+3
View File
@@ -93,6 +93,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
if cfg.AsyncWorkerHardLimit == 0 {
cfg.AsyncWorkerHardLimit = 2048
}
if cfg.AsyncWorkerInstanceHardLimit == 0 {
cfg.AsyncWorkerInstanceHardLimit = 32
}
if cfg.AsyncWorkerRefreshIntervalSeconds == 0 {
cfg.AsyncWorkerRefreshIntervalSeconds = 5
}
@@ -625,4 +625,17 @@ WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).Str
if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
}
first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 3,
})
if err != nil || first.Allocated != 3 || first.GlobalAllocated != 3 {
t.Fatalf("bounded single-worker allocation = %+v, err=%v", first, err)
}
second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 2,
})
if err != nil || second.Allocated != 2 || second.GlobalAllocated != 5 || second.ActiveInstances != 2 {
t.Fatalf("bounded two-worker allocation = %+v, err=%v", second, err)
}
}
+60 -18
View File
@@ -18,6 +18,7 @@ type WorkerRegistrationInput struct {
Site string
Revision string
DesiredCapacity int
CapacityLimit int
HeartbeatStaleAfter time.Duration
}
@@ -25,10 +26,16 @@ type WorkerAllocation struct {
InstanceID string
DesiredCapacity int
Allocated int
GlobalAllocated int
ActiveInstances int
HeartbeatAt time.Time
}
type activeWorkerCapacity struct {
InstanceID string
CapacityLimit int
}
// ActiveWorkerCapacity returns the cluster-wide execution capacity currently
// owned by live Worker instances.
func (s *Store) ActiveWorkerCapacity(ctx context.Context) (int, error) {
@@ -49,6 +56,12 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
if input.DesiredCapacity < 0 {
return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
}
if input.CapacityLimit < 0 {
return WorkerAllocation{}, errors.New("worker capacity limit cannot be negative")
}
if input.CapacityLimit == 0 {
input.CapacityLimit = input.DesiredCapacity
}
staleAfter := input.HeartbeatStaleAfter
if staleAfter < workerHeartbeatStaleAfter {
staleAfter = workerHeartbeatStaleAfter
@@ -70,9 +83,9 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_worker_instances (
instance_id, pod_uid, pod_name, site, revision, status,
desired_capacity, allocated_capacity, started_at, heartbeat_at, updated_at
desired_capacity, capacity_limit, allocated_capacity, started_at, heartbeat_at, updated_at
)
VALUES ($1, $2, $3, $4, $5, 'active', $6, 0, now(), now(), now())
VALUES ($1, $2, $3, $4, $5, 'active', $6, $7, 0, now(), now(), now())
ON CONFLICT (instance_id) DO UPDATE
SET pod_uid = EXCLUDED.pod_uid,
pod_name = EXCLUDED.pod_name,
@@ -80,6 +93,7 @@ SET pod_uid = EXCLUDED.pod_uid,
revision = EXCLUDED.revision,
status = 'active',
desired_capacity = EXCLUDED.desired_capacity,
capacity_limit = EXCLUDED.capacity_limit,
heartbeat_at = now(),
updated_at = now()`,
input.InstanceID,
@@ -88,6 +102,7 @@ SET pod_uid = EXCLUDED.pod_uid,
strings.TrimSpace(input.Site),
strings.TrimSpace(input.Revision),
input.DesiredCapacity,
input.CapacityLimit,
); err != nil {
return WorkerAllocation{}, err
}
@@ -101,7 +116,7 @@ WHERE status <> 'active'
}
rows, err := tx.Query(ctx, `
SELECT instance_id
SELECT instance_id, capacity_limit
FROM gateway_worker_instances
WHERE status = 'active'
AND heartbeat_at > now() - $1::interval
@@ -110,41 +125,37 @@ FOR UPDATE`, staleAfter.String())
if err != nil {
return WorkerAllocation{}, err
}
activeIDs := make([]string, 0)
activeWorkers := make([]activeWorkerCapacity, 0)
for rows.Next() {
var instanceID string
if err := rows.Scan(&instanceID); err != nil {
var worker activeWorkerCapacity
if err := rows.Scan(&worker.InstanceID, &worker.CapacityLimit); err != nil {
rows.Close()
return WorkerAllocation{}, err
}
activeIDs = append(activeIDs, instanceID)
activeWorkers = append(activeWorkers, worker)
}
if err := rows.Err(); err != nil {
rows.Close()
return WorkerAllocation{}, err
}
rows.Close()
if len(activeIDs) == 0 {
if len(activeWorkers) == 0 {
return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat")
}
base := input.DesiredCapacity / len(activeIDs)
remainder := input.DesiredCapacity % len(activeIDs)
allocations, globalAllocated := allocateWorkerCapacities(activeWorkers, input.DesiredCapacity)
allocated := 0
for index, instanceID := range activeIDs {
capacity := base
if index < remainder {
capacity++
}
for _, worker := range activeWorkers {
capacity := allocations[worker.InstanceID]
if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances
SET desired_capacity = $2,
allocated_capacity = $3,
updated_at = now()
WHERE instance_id = $1`, instanceID, input.DesiredCapacity, capacity); err != nil {
WHERE instance_id = $1`, worker.InstanceID, input.DesiredCapacity, capacity); err != nil {
return WorkerAllocation{}, err
}
if instanceID == input.InstanceID {
if worker.InstanceID == input.InstanceID {
allocated = capacity
}
}
@@ -162,11 +173,42 @@ WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
InstanceID: input.InstanceID,
DesiredCapacity: input.DesiredCapacity,
Allocated: allocated,
ActiveInstances: len(activeIDs),
GlobalAllocated: globalAllocated,
ActiveInstances: len(activeWorkers),
HeartbeatAt: heartbeatAt,
}, nil
}
func allocateWorkerCapacities(workers []activeWorkerCapacity, desired int) (map[string]int, int) {
allocations := make(map[string]int, len(workers))
if desired <= 0 || len(workers) == 0 {
return allocations, 0
}
remaining := desired
for remaining > 0 {
progressed := false
for _, worker := range workers {
limit := worker.CapacityLimit
if limit <= 0 {
limit = desired
}
if allocations[worker.InstanceID] >= limit {
continue
}
allocations[worker.InstanceID]++
remaining--
progressed = true
if remaining == 0 {
break
}
}
if !progressed {
break
}
}
return allocations, desired - remaining
}
func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error {
result, err := s.pool.Exec(ctx, `
UPDATE gateway_worker_instances
@@ -0,0 +1,44 @@
package store
import "testing"
func TestAllocateWorkerCapacitiesHonorsInstanceLimits(t *testing.T) {
workers := []activeWorkerCapacity{
{InstanceID: "worker-a", CapacityLimit: 24},
{InstanceID: "worker-b", CapacityLimit: 24},
}
allocations, global := allocateWorkerCapacities(workers, 310)
if global != 48 || allocations["worker-a"] != 24 || allocations["worker-b"] != 24 {
t.Fatalf("allocations=%v global=%d, want 24/24 and 48", allocations, global)
}
}
func TestAllocateWorkerCapacitiesBalancesBelowLimits(t *testing.T) {
workers := []activeWorkerCapacity{
{InstanceID: "worker-a", CapacityLimit: 24},
{InstanceID: "worker-b", CapacityLimit: 24},
}
allocations, global := allocateWorkerCapacities(workers, 5)
if global != 5 || allocations["worker-a"] != 3 || allocations["worker-b"] != 2 {
t.Fatalf("allocations=%v global=%d, want 3/2 and 5", allocations, global)
}
}
func TestAllocateWorkerCapacitiesKeepsSurvivorBounded(t *testing.T) {
workers := []activeWorkerCapacity{{InstanceID: "worker-a", CapacityLimit: 24}}
allocations, global := allocateWorkerCapacities(workers, 310)
if global != 24 || allocations["worker-a"] != 24 {
t.Fatalf("allocations=%v global=%d, want survivor capped at 24", allocations, global)
}
}
func TestAllocateWorkerCapacitiesSupportsUnequalLimits(t *testing.T) {
workers := []activeWorkerCapacity{
{InstanceID: "worker-a", CapacityLimit: 1},
{InstanceID: "worker-b", CapacityLimit: 4},
}
allocations, global := allocateWorkerCapacities(workers, 5)
if global != 5 || allocations["worker-a"] != 1 || allocations["worker-b"] != 4 {
t.Fatalf("allocations=%v global=%d, want 1/4 and 5", allocations, global)
}
}
@@ -0,0 +1,8 @@
ALTER TABLE gateway_worker_instances
ADD COLUMN capacity_limit integer NOT NULL DEFAULT 0,
ADD CONSTRAINT gateway_worker_instances_capacity_limit_check
CHECK (capacity_limit >= 0);
UPDATE gateway_worker_instances
SET capacity_limit = GREATEST(allocated_capacity, 1)
WHERE capacity_limit = 0;
@@ -35,5 +35,5 @@ data:
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES: "268435456"
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES: "536870912"
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED: "true"
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "16"
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "2048"
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS: "5"
@@ -312,8 +312,10 @@ spec:
value: "true"
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "24"
- name: AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT
value: "24"
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
value: "8"
value: "24"
- name: POD_NAMESPACE
valueFrom:
fieldRef:
@@ -439,8 +441,10 @@ spec:
value: "true"
- name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "24"
- name: AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT
value: "24"
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
value: "8"
value: "24"
- name: POD_NAMESPACE
valueFrom:
fieldRef:
+7
View File
@@ -8,6 +8,13 @@
"non-null column addition"
],
"reason": "The migration replaces two existing constraints and one index with expanded definitions. Production preflight confirmed six schema-v1 identity revisions, no OIDC sessions or revocation watermarks, and small affected relations; the non-null columns use constant defaults."
},
"apps/api/migrations/0093_worker_instance_capacity_limit.sql": {
"sha256": "4b6a39ef90e7032555739ccf0568bf7ae0bce23ef98c3bee9c2c8e3bfa13291e",
"allowedViolations": [
"non-null column addition"
],
"reason": "The table stores ephemeral Worker heartbeats and currently has only two production rows. PostgreSQL adds the integer column with a constant default, then the migration backfills each live row from its existing allocation so a rolling deployment cannot assign unsafe failover capacity."
}
}
}