wangbo
ea58d21d03
将同步与异步非文本任务的准入唤醒改为按任务和全局队首推进，批量续租等待者，避免大量请求同时争抢 advisory lock 和数据库连接。 对 Base64 请求解析、素材解码、上游媒体执行和结果物化增加分层并发限制，复用请求素材并释放重复 Gemini wire 数据；生产 API 默认入口解析并发 16，媒体物化并发 8。 新增 1000 个同步 Gemini 图像编辑请求的模拟上游压力验收，10 MiB 输入和输出下全部成功，Heap 峰值增长约 2.58 GiB，并验证共享上传哈希、单次 attempt 与本地零落盘。 验证：go test ./... -count=1；go vet ./...；gofmt -l 无输出；kubectl kustomize deploy/kubernetes/production；10 MiB Gemini Base64 千任务压力测试通过。
295 lines
10 KiB
Go
295 lines
10 KiB
Go
package store
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestNormalizeAndValidateQueuePolicy(t *testing.T) {
|
|
t.Run("defaults maximum wait and preserves extensions", func(t *testing.T) {
|
|
policy, err := NormalizeAndValidateRateLimitPolicy(map[string]any{
|
|
"strategy": "strict",
|
|
"rules": []any{map[string]any{
|
|
"metric": "queue_size",
|
|
"limit": 12,
|
|
"source": "admin",
|
|
}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NormalizeAndValidateRateLimitPolicy() error = %v", err)
|
|
}
|
|
queue, ok := QueueRuleFromPolicy(policy)
|
|
if !ok || queue.Limit != 12 || queue.MaxWaitSeconds != 600 {
|
|
t.Fatalf("queue rule = %+v, %v", queue, ok)
|
|
}
|
|
rules := policy["rules"].([]any)
|
|
rule := rules[0].(map[string]any)
|
|
if rule["source"] != "admin" {
|
|
t.Fatalf("unknown extension was lost: %+v", rule)
|
|
}
|
|
})
|
|
|
|
t.Run("normalizes zero to a removed queue rule", func(t *testing.T) {
|
|
policy, err := NormalizeAndValidateRateLimitPolicy(map[string]any{
|
|
"rules": []any{
|
|
map[string]any{"metric": "rpm", "limit": 30},
|
|
map[string]any{"metric": "queue_size", "limit": 0},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NormalizeAndValidateRateLimitPolicy() error = %v", err)
|
|
}
|
|
if _, ok := QueueRuleFromPolicy(policy); ok {
|
|
t.Fatalf("zero queue rule remained enabled: %+v", policy)
|
|
}
|
|
rules := policy["rules"].([]any)
|
|
if len(rules) != 1 {
|
|
t.Fatalf("rules length = %d, want 1", len(rules))
|
|
}
|
|
})
|
|
|
|
for _, test := range []struct {
|
|
name string
|
|
policy map[string]any
|
|
}{
|
|
{name: "fractional queue", policy: map[string]any{"rules": []any{map[string]any{"metric": "queue_size", "limit": 1.5}}}},
|
|
{name: "queue too large", policy: map[string]any{"rules": []any{map[string]any{"metric": "queue_size", "limit": 10001}}}},
|
|
{name: "wait too large", policy: map[string]any{"rules": []any{map[string]any{"metric": "queue_size", "limit": 1, "maxWaitSeconds": 3601}}}},
|
|
{name: "wait on rpm", policy: map[string]any{"rules": []any{map[string]any{"metric": "rpm", "limit": 1, "maxWaitSeconds": 60}}}},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
if _, err := NormalizeAndValidateRateLimitPolicy(test.policy); err == nil {
|
|
t.Fatalf("invalid policy was accepted: %+v", test.policy)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestEffectiveRateLimitPolicyPrecedence(t *testing.T) {
|
|
policy := func(limit float64) map[string]any {
|
|
return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
input EffectiveRateLimitPolicyInput
|
|
want float64
|
|
ok bool
|
|
}{
|
|
{name: "base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2)}, want: 2, ok: true},
|
|
{name: "platform", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: policy(4)}, want: 4, ok: true},
|
|
{name: "empty platform inherits base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: map[string]any{"rules": []any{}}}, want: 2, ok: true},
|
|
{name: "runtime", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicy: policy(8), RuntimePolicyExplicit: true}, want: 8, ok: true},
|
|
{name: "runtime override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicyOverride: map[string]any{"rateLimitPolicy": policy(16)}}, want: 16, ok: true},
|
|
{name: "model override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "override"}, want: 32, ok: true},
|
|
{name: "model explicit unlimited", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: map[string]any{}, ModelPolicyMode: "override"}, ok: false},
|
|
{name: "model inherit", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "inherit"}, want: 4, ok: true},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got, ok := ConcurrentPolicyCapacity(EffectiveRateLimitPolicy(tt.input))
|
|
if ok != tt.ok || (ok && got != int(tt.want)) {
|
|
t.Fatalf("capacity = (%d, %v), want (%d, %v)", got, ok, int(tt.want), tt.ok)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNormalizeRateLimitPolicyLegacyShapes(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
policy map[string]any
|
|
metric string
|
|
want float64
|
|
}{
|
|
{
|
|
name: "platform concurrent",
|
|
policy: map[string]any{"platformLimits": map[string]any{"max_concurrent_requests": 5.0}},
|
|
metric: "concurrent",
|
|
want: 5,
|
|
},
|
|
{
|
|
name: "model concurrent camel case",
|
|
policy: map[string]any{"modelLimits": map[string]any{"maxConcurrentRequests": 10.0}},
|
|
metric: "concurrent",
|
|
want: 10,
|
|
},
|
|
{
|
|
name: "stricter duplicate wins",
|
|
policy: map[string]any{
|
|
"platformLimits": map[string]any{"max_concurrent_requests": 8.0},
|
|
"modelLimits": map[string]any{"max_concurrent_requests": 3.0},
|
|
},
|
|
metric: "concurrent",
|
|
want: 3,
|
|
},
|
|
{
|
|
name: "requests per minute",
|
|
policy: map[string]any{"model_limits": map[string]any{"max_request_per_minute": 60.0}},
|
|
metric: "rpm",
|
|
want: 60,
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got, ok := RateLimitPolicyMetric(tt.policy, tt.metric)
|
|
if !ok || got != tt.want {
|
|
t.Fatalf("metric %s = (%v, %v), want (%v, true)", tt.metric, got, ok, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConcurrentPolicyCapacityRoundsDown(t *testing.T) {
|
|
policy := map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": 96.9}}}
|
|
got, ok := ConcurrentPolicyCapacity(policy)
|
|
if !ok || got != 96 {
|
|
t.Fatalf("capacity = (%d, %v), want (96, true)", got, ok)
|
|
}
|
|
}
|
|
|
|
func TestAsyncWorkerCapacityAggregation(t *testing.T) {
|
|
policy := func(limit float64) map[string]any {
|
|
return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
policies []map[string]any
|
|
hardLimit int
|
|
wantCapacity int
|
|
wantDesired int
|
|
wantCapped bool
|
|
}{
|
|
{name: "no enabled models", hardLimit: 2048, wantCapacity: 1, wantDesired: 1},
|
|
{name: "finite sum", policies: []map[string]any{policy(64), policy(32)}, hardLimit: 2048, wantCapacity: 96, wantDesired: 96},
|
|
{name: "unlimited model", policies: []map[string]any{policy(64), {}}, hardLimit: 2048, wantCapacity: 2048, wantDesired: 2048},
|
|
{name: "hard limit cap", policies: []map[string]any{policy(80), policy(80)}, hardLimit: 96, wantCapacity: 96, wantDesired: 160, wantCapped: true},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := asyncWorkerCapacityFromPolicies(tt.policies, tt.hardLimit)
|
|
if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped {
|
|
t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAsyncWorkerCapacityRespectsUserGroupCeiling(t *testing.T) {
|
|
policy := func(limit float64) map[string]any {
|
|
return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
models []map[string]any
|
|
groups []map[string]any
|
|
hardLimit int
|
|
wantCapacity int
|
|
wantDesired int
|
|
wantCapped bool
|
|
}{
|
|
{
|
|
name: "group ceiling prevents worker oversubscription",
|
|
models: []map[string]any{{}},
|
|
groups: []map[string]any{policy(3), policy(10)},
|
|
hardLimit: 2048,
|
|
wantCapacity: 13,
|
|
wantDesired: 13,
|
|
},
|
|
{
|
|
name: "model ceiling is stricter",
|
|
models: []map[string]any{policy(5), policy(7)},
|
|
groups: []map[string]any{policy(300)},
|
|
hardLimit: 2048,
|
|
wantCapacity: 12,
|
|
wantDesired: 12,
|
|
},
|
|
{
|
|
name: "both policy sets unlimited use hard limit",
|
|
models: []map[string]any{{}},
|
|
groups: []map[string]any{{}},
|
|
hardLimit: 256,
|
|
wantCapacity: 256,
|
|
wantDesired: 256,
|
|
},
|
|
{
|
|
name: "finite group desired still reports hard cap",
|
|
models: []map[string]any{{}},
|
|
groups: []map[string]any{policy(500)},
|
|
hardLimit: 256,
|
|
wantCapacity: 256,
|
|
wantDesired: 500,
|
|
wantCapped: true,
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := asyncWorkerCapacityFromPolicySets(tt.models, tt.groups, tt.hardLimit)
|
|
if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped {
|
|
t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAdmissionDeadlineSeparatesExecutionSlotWaitFromPolicyQueueTimeout(t *testing.T) {
|
|
now := time.Now()
|
|
infrastructureDeadline := admissionDeadline([]AdmissionScope{{
|
|
ConcurrentLimit: 10,
|
|
}})
|
|
if wait := infrastructureDeadline.Sub(now); wait < 23*time.Hour || wait > executionSlotWaitMax+time.Second {
|
|
t.Fatalf("execution-slot deadline wait=%s, want about %s", wait, executionSlotWaitMax)
|
|
}
|
|
|
|
defaultPolicyDeadline := admissionDeadline([]AdmissionScope{{
|
|
ConcurrentLimit: 10,
|
|
QueueLimit: 20,
|
|
}})
|
|
if wait := defaultPolicyDeadline.Sub(now); wait < 599*time.Second || wait > 601*time.Second {
|
|
t.Fatalf("default policy deadline wait=%s, want about 600s", wait)
|
|
}
|
|
|
|
explicitPolicyDeadline := admissionDeadline([]AdmissionScope{{
|
|
ConcurrentLimit: 10,
|
|
QueueLimit: 20,
|
|
MaxWaitSeconds: 90,
|
|
}})
|
|
if wait := explicitPolicyDeadline.Sub(now); wait < 89*time.Second || wait > 91*time.Second {
|
|
t.Fatalf("explicit policy deadline wait=%s, want about 90s", wait)
|
|
}
|
|
}
|
|
|
|
func TestQueuedAdmissionCountsFIFOWaitersAsVirtualCapacity(t *testing.T) {
|
|
scope := AdmissionScope{
|
|
ScopeType: "platform_model",
|
|
ScopeKey: "model-1",
|
|
ConcurrentLimit: 1,
|
|
Amount: 1,
|
|
QueueLimit: 2,
|
|
}
|
|
for waiting := 0; waiting <= 2; waiting++ {
|
|
if err := validateQueuedAdmissionCapacity([]admissionScopeState{{
|
|
Scope: scope,
|
|
Waiting: waiting,
|
|
}}); err != nil {
|
|
t.Fatalf("waiting=%d rejected within concurrent=1 queue=2: %v", waiting, err)
|
|
}
|
|
}
|
|
err := validateQueuedAdmissionCapacity([]admissionScopeState{{
|
|
Scope: scope,
|
|
Waiting: 3,
|
|
}})
|
|
var limitErr *RateLimitExceededError
|
|
if !errors.As(err, &limitErr) || limitErr.Reason != "queue_full" {
|
|
t.Fatalf("fourth task error=%v, want queue_full", err)
|
|
}
|
|
|
|
scope.QueueLimit = 0
|
|
if err := validateQueuedAdmissionCapacity([]admissionScopeState{{
|
|
Scope: scope,
|
|
Waiting: 1,
|
|
}}); !errors.As(err, &limitErr) || limitErr.Metric != "concurrent" {
|
|
t.Fatalf("queue-disabled second task error=%v, want concurrent rate limit", err)
|
|
}
|
|
}
|