perf(queue): 按策略动态扩缩异步 Worker

将 River 执行容量按平台模型与用户组的有效并发策略动态调整,并为长任务续租、并发租约原子抢占和限流退避补充保护。\n\n统一平台模型限流继承语义,兼容历史 platformLimits/modelLimits,并为三个迁移图像模型建立独立并发租约。补充管理端显式继承/覆盖配置、指标、单元测试及隔离 PostgreSQL 验收。\n\n验证:go test ./... -count=1;pnpm lint;pnpm test;pnpm build;隔离 PostgreSQL 并发原子性/续租测试;128 任务与三分钟长任务动态 Worker 验收。
This commit is contained in:
2026-07-24 12:26:56 +08:00
parent 290b8c1854
commit 6c5daf29ca
37 changed files with 2365 additions and 205 deletions
@@ -0,0 +1,152 @@
package store
import (
"context"
"fmt"
)
type AsyncWorkerCapacitySnapshot struct {
Capacity int
Desired int
HardLimit int
Capped bool
EnabledModels int
UnlimitedModels int
EnabledGroups int
UnlimitedGroups int
ModelDesired int
GroupDesired int
}
func (s *Store) AsyncWorkerCapacity(ctx context.Context, hardLimit int) (AsyncWorkerCapacitySnapshot, error) {
if hardLimit < 1 {
return AsyncWorkerCapacitySnapshot{}, fmt.Errorf("async worker hard limit must be positive")
}
rows, err := s.pool.Query(ctx, `
SELECT COALESCE(b.default_rate_limit_policy, '{}'::jsonb),
p.rate_limit_policy,
COALESCE(rp.rate_limit_policy, '{}'::jsonb),
(m.runtime_policy_set_id IS NOT NULL),
COALESCE(m.runtime_policy_override, '{}'::jsonb),
m.rate_limit_policy,
m.rate_limit_policy_mode
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
WHERE p.status = 'enabled'
AND p.deleted_at IS NULL
AND m.enabled = true`)
if err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
defer rows.Close()
modelPolicies := make([]map[string]any, 0)
for rows.Next() {
var basePolicyBytes, platformPolicyBytes, runtimePolicyBytes []byte
var runtimeOverrideBytes, modelPolicyBytes []byte
var runtimeExplicit bool
var modelPolicyMode string
if err := rows.Scan(
&basePolicyBytes,
&platformPolicyBytes,
&runtimePolicyBytes,
&runtimeExplicit,
&runtimeOverrideBytes,
&modelPolicyBytes,
&modelPolicyMode,
); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
modelPolicies = append(modelPolicies, EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
BasePolicy: decodeObject(basePolicyBytes),
PlatformPolicy: decodeObject(platformPolicyBytes),
RuntimePolicy: decodeObject(runtimePolicyBytes),
RuntimePolicyExplicit: runtimeExplicit,
RuntimePolicyOverride: decodeObject(runtimeOverrideBytes),
ModelPolicy: decodeObject(modelPolicyBytes),
ModelPolicyMode: modelPolicyMode,
}))
}
if err := rows.Err(); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
groupRows, err := s.pool.Query(ctx, `
SELECT rate_limit_policy
FROM gateway_user_groups
WHERE status = 'active'`)
if err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
defer groupRows.Close()
groupPolicies := make([]map[string]any, 0)
for groupRows.Next() {
var policyBytes []byte
if err := groupRows.Scan(&policyBytes); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
groupPolicies = append(groupPolicies, NormalizeRateLimitPolicy(decodeObject(policyBytes)))
}
if err := groupRows.Err(); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
return asyncWorkerCapacityFromPolicySets(modelPolicies, groupPolicies, hardLimit), nil
}
func asyncWorkerCapacityFromPolicies(policies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot {
return asyncWorkerCapacityFromPolicySets(policies, nil, hardLimit)
}
func asyncWorkerCapacityFromPolicySets(modelPolicies []map[string]any, groupPolicies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot {
snapshot := AsyncWorkerCapacitySnapshot{
HardLimit: hardLimit,
EnabledModels: len(modelPolicies),
EnabledGroups: len(groupPolicies),
}
modelDesired, modelFinite, unlimitedModels := concurrentPolicySetCapacity(modelPolicies)
groupDesired, groupFinite, unlimitedGroups := concurrentPolicySetCapacity(groupPolicies)
snapshot.ModelDesired = modelDesired
snapshot.GroupDesired = groupDesired
snapshot.UnlimitedModels = unlimitedModels
snapshot.UnlimitedGroups = unlimitedGroups
desired := hardLimit
switch {
case snapshot.EnabledModels == 0:
desired = 1
case modelFinite && groupFinite:
desired = min(modelDesired, groupDesired)
case modelFinite:
desired = modelDesired
case groupFinite:
desired = groupDesired
}
if desired < 1 {
desired = 1
}
snapshot.Desired = desired
snapshot.Capacity = desired
if snapshot.Capacity > hardLimit {
snapshot.Capacity = hardLimit
snapshot.Capped = true
}
return snapshot
}
func concurrentPolicySetCapacity(policies []map[string]any) (total int, finite bool, unlimited int) {
if len(policies) == 0 {
return 0, false, 0
}
finite = true
for _, policy := range policies {
capacity, policyFinite := ConcurrentPolicyCapacity(policy)
if !policyFinite {
unlimited++
finite = false
continue
}
total += capacity
}
return total, finite, unlimited
}
+22 -3
View File
@@ -25,7 +25,7 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
COALESCE(p.base_url, ''),
p.auth_type, p.credentials, p.config, p.default_pricing_mode,
p.default_discount_factor::float8, COALESCE(p.pricing_rule_set_id::text, ''),
p.retry_policy, p.rate_limit_policy,
p.retry_policy, p.rate_limit_policy, COALESCE(b.default_rate_limit_policy, '{}'::jsonb),
COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''),
@@ -54,8 +54,11 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
COALESCE(b.pricing_rule_set_id::text, ''),
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
m.permission_config, m.retry_policy, m.rate_limit_policy, m.rate_limit_policy_mode,
COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
(m.runtime_policy_set_id IS NOT NULL),
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb),
COALESCE(m.runtime_policy_override, '{}'::jsonb),
COALESCE(rp.retry_policy, '{}'::jsonb), COALESCE(rp.rate_limit_policy, '{}'::jsonb),
COALESCE(rp.auto_disable_policy, '{}'::jsonb), COALESCE(rp.degrade_policy, '{}'::jsonb),
COALESCE(con.active, 0)::float8,
@@ -177,6 +180,7 @@ WHERE p.status = 'enabled'
var platformConfig []byte
var platformRetryPolicy []byte
var platformRateLimitPolicy []byte
var baseRateLimitPolicy []byte
var capabilities []byte
var capabilityOverride []byte
var baseBilling []byte
@@ -186,6 +190,7 @@ WHERE p.status = 'enabled'
var modelRetryPolicy []byte
var modelRateLimitPolicy []byte
var runtimePolicyOverride []byte
var rateLimitRuntimeOverride []byte
var runtimeRetryPolicy []byte
var runtimeRateLimitPolicy []byte
var autoDisablePolicy []byte
@@ -222,6 +227,7 @@ WHERE p.status = 'enabled'
&item.PlatformPricingRuleSetID,
&platformRetryPolicy,
&platformRateLimitPolicy,
&baseRateLimitPolicy,
&item.PlatformPriority,
&item.PlatformModelID,
&item.BaseModelID,
@@ -244,8 +250,11 @@ WHERE p.status = 'enabled'
&permissionConfig,
&modelRetryPolicy,
&modelRateLimitPolicy,
&item.ModelRateLimitPolicyMode,
&item.RuntimePolicySetID,
&item.RuntimePolicyExplicit,
&runtimePolicyOverride,
&rateLimitRuntimeOverride,
&runtimeRetryPolicy,
&runtimeRateLimitPolicy,
&autoDisablePolicy,
@@ -274,6 +283,7 @@ WHERE p.status = 'enabled'
item.PlatformConfig = decodeObject(platformConfig)
item.PlatformRetryPolicy = decodeObject(platformRetryPolicy)
item.PlatformRateLimitPolicy = decodeObject(platformRateLimitPolicy)
item.BaseRateLimitPolicy = decodeObject(baseRateLimitPolicy)
item.Capabilities = decodeObject(capabilities)
item.CapabilityOverride = decodeObject(capabilityOverride)
item.BaseBillingConfig = decodeObject(baseBilling)
@@ -283,6 +293,7 @@ WHERE p.status = 'enabled'
item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
item.RateLimitRuntimeOverride = decodeObject(rateLimitRuntimeOverride)
item.RuntimeRetryPolicy = decodeObject(runtimeRetryPolicy)
item.RuntimeRateLimitPolicy = decodeObject(runtimeRateLimitPolicy)
item.AutoDisablePolicy = decodeObject(autoDisablePolicy)
@@ -303,7 +314,15 @@ WHERE p.status = 'enabled'
LastObservedUnix: cacheLastObservedUnix,
})
applyRuntimeCandidateLoad(&item, runtimeCandidateLoadInput{
Policy: effectiveModelRateLimitPolicy(item.PlatformRateLimitPolicy, item.RuntimeRateLimitPolicy, item.RuntimePolicySetID, item.RuntimePolicyOverride, item.ModelRateLimitPolicy),
Policy: EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
BasePolicy: item.BaseRateLimitPolicy,
PlatformPolicy: item.PlatformRateLimitPolicy,
RuntimePolicy: item.RuntimeRateLimitPolicy,
RuntimePolicyExplicit: item.RuntimePolicyExplicit,
RuntimePolicyOverride: item.RateLimitRuntimeOverride,
ModelPolicy: item.ModelRateLimitPolicy,
ModelPolicyMode: item.ModelRateLimitPolicyMode,
}),
ConcurrentActive: concurrentActive,
QueuedWaiting: queuedWaiting,
RPMUsed: rpmUsed,
+10 -12
View File
@@ -143,18 +143,13 @@ SELECT EXISTS (
// soon as the base pricing rule changes and can mask the authoritative rule.
billingConfig := input.BillingConfig
explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
rateLimitPolicyMode := NormalizeRateLimitPolicyMode(input.RateLimitPolicyMode, input.RateLimitPolicy != nil)
rateLimitPolicy := input.RateLimitPolicy
if len(rateLimitPolicy) == 0 && explicitRuntimePolicySetID == "" {
rateLimitPolicy = base.DefaultRateLimitPolicy
if rateLimitPolicyMode == RateLimitPolicyModeInherit {
rateLimitPolicy = nil
}
runtimePolicySetID := explicitRuntimePolicySetID
if runtimePolicySetID == "" {
runtimePolicySetID = base.RuntimePolicySetID
}
runtimePolicyOverride := input.RuntimePolicyOverride
if len(runtimePolicyOverride) == 0 {
runtimePolicyOverride = base.RuntimePolicyOverride
}
capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
@@ -189,14 +184,14 @@ SELECT EXISTS (
INSERT INTO platform_models (
platform_id, base_model_id, model_name, provider_model_name, model_alias, model_type, display_name,
capability_override, capabilities, pricing_mode, discount_factor,
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy,
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode,
runtime_policy_set_id, runtime_policy_override, enabled
)
VALUES (
$1::uuid, $2::uuid, $3, NULLIF($4, ''), NULLIF($5, ''), $6::jsonb, $7,
$8::jsonb, $9::jsonb, $10, $11::numeric,
NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb,
NULLIF($18, '')::uuid, $19::jsonb, true
NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb, $18,
NULLIF($19, '')::uuid, $20::jsonb, true
)
ON CONFLICT (platform_id, model_name) DO UPDATE
SET base_model_id = EXCLUDED.base_model_id,
@@ -213,6 +208,7 @@ SET base_model_id = EXCLUDED.base_model_id,
permission_config = EXCLUDED.permission_config,
retry_policy = EXCLUDED.retry_policy,
rate_limit_policy = EXCLUDED.rate_limit_policy,
rate_limit_policy_mode = EXCLUDED.rate_limit_policy_mode,
runtime_policy_set_id = EXCLUDED.runtime_policy_set_id,
runtime_policy_override = EXCLUDED.runtime_policy_override,
enabled = true,
@@ -221,7 +217,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
COALESCE(NULLIF(provider_model_name, ''), model_name), COALESCE(model_alias, ''), model_type, display_name, capability_override,
capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config,
permission_config, retry_policy, rate_limit_policy, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
enabled, created_at, updated_at`,
input.PlatformID,
@@ -241,6 +237,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
string(permissionJSON),
string(retryJSON),
string(rateLimitJSON),
rateLimitPolicyMode,
runtimePolicySetID,
string(runtimePolicyOverrideJSON),
).Scan(
@@ -262,6 +259,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
&permissionBytes,
&retryPolicyBytes,
&rateLimitPolicyBytes,
&model.RateLimitPolicyMode,
&model.RuntimePolicySetID,
&runtimePolicyOverrideBytes,
&model.CooldownUntil,
+12 -1
View File
@@ -235,6 +235,8 @@ type PlatformModel struct {
Capabilities map[string]any `json:"capabilities,omitempty"`
BaseCapabilities map[string]any `json:"-"`
BaseBillingConfig map[string]any `json:"-"`
BaseRateLimitPolicy map[string]any `json:"-"`
BaseRuntimePolicySetID string `json:"-"`
BasePricingRuleSetID string `json:"-"`
PlatformPricingRuleSetID string `json:"-"`
PricingMode string `json:"pricingMode"`
@@ -245,6 +247,7 @@ type PlatformModel struct {
PermissionConfig map[string]any `json:"permissionConfig,omitempty"`
RetryPolicy map[string]any `json:"retryPolicy,omitempty"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
RateLimitPolicyMode string `json:"rateLimitPolicyMode" enums:"inherit,override"`
RuntimePolicySetID string `json:"runtimePolicySetId,omitempty"`
RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride,omitempty"`
CooldownUntil string `json:"cooldownUntil,omitempty"`
@@ -956,7 +959,9 @@ SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.p
COALESCE(b.base_billing_config, '{}'::jsonb), COALESCE(b.pricing_rule_set_id::text, ''),
COALESCE(p.pricing_rule_set_id::text, ''), m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
m.permission_config, m.retry_policy, m.rate_limit_policy, m.rate_limit_policy_mode,
COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
COALESCE(b.default_rate_limit_policy, '{}'::jsonb), COALESCE(b.runtime_policy_set_id::text, ''),
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
m.enabled, m.created_at, m.updated_at
FROM platform_models m
@@ -964,6 +969,7 @@ JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN LATERAL (
SELECT catalog.invocation_name, catalog.display_name, catalog.model_type,
catalog.capabilities, catalog.base_billing_config, catalog.pricing_rule_set_id,
catalog.default_rate_limit_policy, catalog.runtime_policy_set_id,
COALESCE((
SELECT jsonb_agg(DISTINCT compatibility_alias.alias ORDER BY compatibility_alias.alias)
FROM model_compatibility_aliases compatibility_alias
@@ -1004,6 +1010,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
var retryPolicy []byte
var rateLimitPolicy []byte
var runtimePolicyOverride []byte
var baseRateLimitPolicy []byte
var modelTypeBytes []byte
var legacyAliasesBytes []byte
if err := rows.Scan(
@@ -1032,8 +1039,11 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
&permissionConfig,
&retryPolicy,
&rateLimitPolicy,
&model.RateLimitPolicyMode,
&model.RuntimePolicySetID,
&runtimePolicyOverride,
&baseRateLimitPolicy,
&model.BaseRuntimePolicySetID,
&model.CooldownUntil,
&model.Enabled,
&model.CreatedAt,
@@ -1053,6 +1063,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
model.RetryPolicy = decodeObject(retryPolicy)
model.RateLimitPolicy = decodeObject(rateLimitPolicy)
model.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
model.BaseRateLimitPolicy = decodeObject(baseRateLimitPolicy)
models = append(models, model)
}
return models, rows.Err()
@@ -0,0 +1,167 @@
package store
import (
"math"
"strings"
)
const (
RateLimitPolicyModeInherit = "inherit"
RateLimitPolicyModeOverride = "override"
)
type EffectiveRateLimitPolicyInput struct {
BasePolicy map[string]any
PlatformPolicy map[string]any
RuntimePolicy map[string]any
RuntimePolicyExplicit bool
RuntimePolicyOverride map[string]any
ModelPolicy map[string]any
ModelPolicyMode string
}
// EffectiveRateLimitPolicy resolves one authoritative policy. Platform
// policies are defaults for every bound model; explicit model/runtime settings
// replace the complete policy instead of merging individual metrics.
func EffectiveRateLimitPolicy(input EffectiveRateLimitPolicyInput) map[string]any {
policy := input.BasePolicy
if policySpecified(input.PlatformPolicy) {
policy = input.PlatformPolicy
}
if input.RuntimePolicyExplicit {
policy = input.RuntimePolicy
}
if raw, ok := input.RuntimePolicyOverride["rateLimitPolicy"]; ok {
policy, _ = raw.(map[string]any)
}
mode := NormalizeRateLimitPolicyMode(input.ModelPolicyMode, input.ModelPolicy != nil)
if mode == RateLimitPolicyModeOverride {
policy = input.ModelPolicy
}
return NormalizeRateLimitPolicy(policy)
}
func NormalizeRateLimitPolicyMode(mode string, policyProvided bool) string {
switch strings.ToLower(strings.TrimSpace(mode)) {
case RateLimitPolicyModeOverride:
return RateLimitPolicyModeOverride
case RateLimitPolicyModeInherit:
return RateLimitPolicyModeInherit
default:
if policyProvided {
return RateLimitPolicyModeOverride
}
return RateLimitPolicyModeInherit
}
}
func RateLimitPolicyMetric(policy map[string]any, metric string) (float64, bool) {
rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any)
for _, rawRule := range rules {
rule, _ := rawRule.(map[string]any)
if strings.TrimSpace(stringValue(rule["metric"])) != metric {
continue
}
value := floatValue(rule["limit"])
return value, true
}
return 0, false
}
// NormalizeRateLimitPolicy keeps the canonical rules contract while accepting
// policies imported from server-main before that contract existed. Runtime
// enforcement and worker sizing must agree on these legacy limits; otherwise a
// configured max_concurrent_requests silently becomes unlimited.
func NormalizeRateLimitPolicy(policy map[string]any) map[string]any {
if policy == nil {
return nil
}
out := clonePolicy(policy)
rules, _ := out["rules"].([]any)
if len(rules) > 0 {
return out
}
legacyScopes := []map[string]any{policy}
for _, key := range []string{"platformLimits", "modelLimits", "platform_limits", "model_limits"} {
if nested, ok := policy[key].(map[string]any); ok {
legacyScopes = append(legacyScopes, nested)
}
}
if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
"max_concurrent_requests", "maxConcurrentRequests", "concurrent"); ok {
rules = append(rules, map[string]any{
"metric": "concurrent",
"limit": limit,
"leaseTtlSeconds": 120,
})
}
if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
"max_request_per_minute", "maxRequestsPerMinute", "rpm"); ok {
rules = append(rules, map[string]any{
"metric": "rpm",
"limit": limit,
"windowSeconds": 60,
})
}
if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
"max_tokens_per_minute", "maxTokensPerMinute", "tpm_total"); ok {
rules = append(rules, map[string]any{
"metric": "tpm_total",
"limit": limit,
"windowSeconds": 60,
})
}
if len(rules) > 0 {
out["rules"] = rules
}
return out
}
func ConcurrentPolicyCapacity(policy map[string]any) (int, bool) {
limit, ok := RateLimitPolicyMetric(policy, "concurrent")
if !ok || limit <= 0 {
return 0, false
}
if limit < 1 {
return 1, true
}
if limit >= float64(math.MaxInt) {
return math.MaxInt, true
}
return int(math.Floor(limit)), true
}
func policySpecified(policy map[string]any) bool {
rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any)
return len(rules) > 0
}
func lowestPositiveLegacyLimit(scopes []map[string]any, keys ...string) (float64, bool) {
limit := 0.0
found := false
for _, scope := range scopes {
for _, key := range keys {
value := floatValue(scope[key])
if value <= 0 {
continue
}
if !found || value < limit {
limit = value
found = true
}
}
}
return limit, found
}
func clonePolicy(policy map[string]any) map[string]any {
if policy == nil {
return nil
}
out := make(map[string]any, len(policy))
for key, value := range policy {
out[key] = value
}
return out
}
@@ -0,0 +1,120 @@
package store
import (
"context"
"fmt"
"net/url"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestRateLimitPolicyModeMigrationClassifiesExistingRows(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run rate-limit migration PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
schema := fmt.Sprintf("rate_limit_mode_%d", time.Now().UnixNano())
adminPool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect migration test database: %v", err)
}
defer adminPool.Close()
if _, err := adminPool.Exec(ctx, `CREATE SCHEMA `+schema); err != nil {
t.Fatalf("create migration test schema: %v", err)
}
defer adminPool.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schema+` CASCADE`)
schemaURL, err := url.Parse(databaseURL)
if err != nil {
t.Fatalf("parse test database url: %v", err)
}
query := schemaURL.Query()
query.Set("search_path", schema)
schemaURL.RawQuery = query.Encode()
pool, err := pgxpool.New(ctx, schemaURL.String())
if err != nil {
t.Fatalf("connect migration test schema: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `
CREATE TABLE base_model_catalog (
id uuid PRIMARY KEY,
default_rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
runtime_policy_set_id uuid,
runtime_policy_override jsonb NOT NULL DEFAULT '{}'::jsonb
);
CREATE TABLE integration_platforms (
id uuid PRIMARY KEY,
rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
updated_at timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE platform_models (
id uuid PRIMARY KEY,
base_model_id uuid,
rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
runtime_policy_set_id uuid,
runtime_policy_override jsonb NOT NULL DEFAULT '{}'::jsonb,
updated_at timestamptz NOT NULL DEFAULT now()
);
INSERT INTO integration_platforms (id, rate_limit_policy)
VALUES ('20000000-0000-0000-0000-000000000001', '{"rules":[]}');
INSERT INTO base_model_catalog (id, default_rate_limit_policy)
VALUES ('00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}');
INSERT INTO platform_models (id, base_model_id, rate_limit_policy, runtime_policy_override) VALUES
('10000000-0000-0000-0000-000000000001', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}', '{}'),
('10000000-0000-0000-0000-000000000002', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":8}]}', '{}'),
('10000000-0000-0000-0000-000000000003', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}', '{"rateLimitPolicy":{"rules":[]}}'),
('10000000-0000-0000-0000-000000000004', NULL, '{}', '{}');`); err != nil {
t.Fatalf("seed pre-migration rows: %v", err)
}
_, currentFile, _, _ := runtime.Caller(0)
migrationPath := filepath.Join(filepath.Dir(currentFile), "..", "..", "migrations", "0080_platform_model_rate_limit_policy_mode.sql")
migration, err := os.ReadFile(migrationPath)
if err != nil {
t.Fatalf("read migration: %v", err)
}
if _, err := pool.Exec(ctx, string(migration)); err != nil {
t.Fatalf("apply migration: %v", err)
}
rows, err := pool.Query(ctx, `SELECT id::text, rate_limit_policy_mode FROM platform_models ORDER BY id`)
if err != nil {
t.Fatalf("read classified rows: %v", err)
}
defer rows.Close()
want := []string{"inherit", "override", "override", "inherit"}
index := 0
for rows.Next() {
var id, mode string
if err := rows.Scan(&id, &mode); err != nil {
t.Fatalf("scan classified row: %v", err)
}
if index >= len(want) || mode != want[index] {
t.Fatalf("row %s mode=%s, want=%s", id, mode, want[index])
}
index++
}
if err := rows.Err(); err != nil {
t.Fatalf("read classified rows: %v", err)
}
if index != len(want) {
t.Fatalf("classified %d rows, want %d", index, len(want))
}
var normalizedPlatformPolicy string
if err := pool.QueryRow(ctx, `
SELECT rate_limit_policy::text
FROM integration_platforms
WHERE id = '20000000-0000-0000-0000-000000000001'`).Scan(&normalizedPlatformPolicy); err != nil {
t.Fatalf("read normalized platform policy: %v", err)
}
if normalizedPlatformPolicy != "{}" {
t.Fatalf("normalized platform policy=%s, want={}", normalizedPlatformPolicy)
}
}
@@ -0,0 +1,169 @@
package store
import "testing"
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)
}
})
}
}
@@ -0,0 +1,19 @@
package store
import (
"testing"
"time"
)
func TestConcurrencyRetryAfterBoundsPolling(t *testing.T) {
if got := concurrencyRetryAfter(time.Time{}); got != 5*time.Second {
t.Fatalf("zero expiry retry=%s, want=5s", got)
}
if got := concurrencyRetryAfter(time.Now().Add(30 * time.Second)); got != 5*time.Second {
t.Fatalf("distant expiry retry=%s, want=5s", got)
}
got := concurrencyRetryAfter(time.Now().Add(2500 * time.Millisecond))
if got < 2*time.Second || got > 2500*time.Millisecond {
t.Fatalf("near expiry retry=%s, want within [2s,2.5s]", got)
}
}
+19 -51
View File
@@ -143,8 +143,10 @@ func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimi
p.priority, p.dynamic_priority, COALESCE(p.dynamic_priority, p.priority),
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
m.model_type, m.display_name, m.enabled,
p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(m.runtime_policy_set_id::text, b.runtime_policy_set_id::text, ''),
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy,
COALESCE(b.default_rate_limit_policy, '{}'::jsonb), p.rate_limit_policy,
COALESCE(rp.rate_limit_policy, '{}'::jsonb),
(m.runtime_policy_set_id IS NOT NULL),
COALESCE(m.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy, m.rate_limit_policy_mode,
COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
COALESCE(con.active, 0)::float8,
@@ -217,11 +219,13 @@ ORDER BY p.priority ASC, m.model_name ASC`)
for rows.Next() {
var item ModelRateLimitStatus
var modelTypeBytes []byte
var basePolicyBytes []byte
var platformPolicyBytes []byte
var runtimePolicyBytes []byte
var runtimePolicySetID string
var runtimePolicyExplicit bool
var runtimeOverrideBytes []byte
var modelPolicyBytes []byte
var modelPolicyMode string
var platformDynamicPriority sql.NullInt64
var platformCooldownUntil string
var modelCooldownUntil string
@@ -248,11 +252,13 @@ ORDER BY p.priority ASC, m.model_name ASC`)
&modelTypeBytes,
&item.DisplayName,
&item.Enabled,
&basePolicyBytes,
&platformPolicyBytes,
&runtimePolicyBytes,
&runtimePolicySetID,
&runtimePolicyExplicit,
&runtimeOverrideBytes,
&modelPolicyBytes,
&modelPolicyMode,
&platformCooldownUntil,
&modelCooldownUntil,
&concurrentCurrent,
@@ -268,13 +274,15 @@ ORDER BY p.priority ASC, m.model_name ASC`)
}
item.PlatformDynamicPriority = intPointerFromNull(platformDynamicPriority)
item.ModelType = decodeStringArray(modelTypeBytes)
policy := effectiveModelRateLimitPolicy(
decodeObject(platformPolicyBytes),
decodeObject(runtimePolicyBytes),
runtimePolicySetID,
decodeObject(runtimeOverrideBytes),
decodeObject(modelPolicyBytes),
)
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
BasePolicy: decodeObject(basePolicyBytes),
PlatformPolicy: decodeObject(platformPolicyBytes),
RuntimePolicy: decodeObject(runtimePolicyBytes),
RuntimePolicyExplicit: runtimePolicyExplicit,
RuntimePolicyOverride: decodeObject(runtimeOverrideBytes),
ModelPolicy: decodeObject(modelPolicyBytes),
ModelPolicyMode: modelPolicyMode,
})
item.PlatformCooldownUntil = platformCooldownUntil
item.ModelCooldownUntil = modelCooldownUntil
item.RateLimitPolicy = policy
@@ -491,46 +499,6 @@ func platformPolicyEventFromPayload(id string, taskID string, eventType string,
}
}
func effectiveModelRateLimitPolicy(platformPolicy map[string]any, runtimePolicy map[string]any, runtimePolicySetID string, runtimeOverride map[string]any, modelPolicy map[string]any) map[string]any {
policy := platformPolicy
if strings.TrimSpace(runtimePolicySetID) != "" {
policy = runtimePolicy
} else if hasRateLimitRules(runtimePolicy) {
policy = shallowMergeMap(policy, runtimePolicy)
}
if _, hasOverride := runtimeOverride["rateLimitPolicy"]; hasOverride {
nested, _ := runtimeOverride["rateLimitPolicy"].(map[string]any)
if len(nested) == 0 {
policy = nil
} else {
policy = shallowMergeMap(policy, nested)
}
}
if hasRateLimitRules(modelPolicy) {
policy = shallowMergeMap(policy, modelPolicy)
}
if hasRateLimitRules(policy) {
return policy
}
return nil
}
func hasRateLimitRules(policy map[string]any) bool {
rules, _ := policy["rules"].([]any)
return len(rules) > 0
}
func shallowMergeMap(base map[string]any, override map[string]any) map[string]any {
out := map[string]any{}
for key, value := range base {
out[key] = value
}
for key, value := range override {
out[key] = value
}
return out
}
func rateLimitForMetric(policy map[string]any, metric string) float64 {
rules, _ := policy["rules"].([]any)
for _, rawRule := range rules {
@@ -6,23 +6,23 @@ import (
)
func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
map[string]any{"rules": []any{
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
PlatformPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
map[string]any{"metric": "tpm_total", "limit": 100000},
}},
map[string]any{"rules": []any{
RuntimePolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 120},
map[string]any{"metric": "tpm_total", "limit": 240000},
map[string]any{"metric": "concurrent", "limit": 6},
}},
"runtime-policy-1",
map[string]any{},
map[string]any{"rules": []any{
RuntimePolicyExplicit: true,
ModelPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 30},
map[string]any{"metric": "concurrent", "limit": 2},
}},
)
ModelPolicyMode: "override",
})
if got := rateLimitForMetric(policy, "rpm"); got != 30 {
t.Fatalf("expected model rpm limit to win, got %v", got)
@@ -36,16 +36,15 @@ func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing
}
func TestEffectiveModelRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
map[string]any{"rules": []any{
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
PlatformPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
map[string]any{"metric": "tpm_total", "limit": 100000},
}},
map[string]any{"rules": []any{}},
"runtime-policy-1",
map[string]any{},
map[string]any{},
)
RuntimePolicy: map[string]any{"rules": []any{}},
RuntimePolicyExplicit: true,
ModelPolicyMode: "inherit",
})
if got := rateLimitForMetric(policy, "rpm"); got != 0 {
t.Fatalf("expected empty runtime policy rpm to mean unlimited, got %v", got)
@@ -56,17 +55,16 @@ func TestEffectiveModelRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *tes
}
func TestEffectiveModelRateLimitPolicyTreatsNegativeLimitAsUnlimited(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
map[string]any{"rules": []any{
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
PlatformPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
}},
map[string]any{"rules": []any{
RuntimePolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": -1},
}},
"runtime-policy-1",
map[string]any{},
map[string]any{},
)
RuntimePolicyExplicit: true,
ModelPolicyMode: "inherit",
})
if got := rateLimitForMetric(policy, "rpm"); got != -1 {
t.Fatalf("expected negative runtime rpm marker to be preserved, got %v", got)
+87 -19
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"sort"
"time"
"github.com/jackc/pgx/v5"
@@ -17,6 +18,8 @@ type RuntimeRecoveryResult struct {
RequeuedAsyncTasks int64 `json:"requeuedAsyncTasks"`
}
var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost")
func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
@@ -24,6 +27,26 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
}
defer tx.Rollback(ctx)
lockKeys := make([]string, 0)
lockKeySet := make(map[string]struct{})
for _, reservation := range reservations {
if reservation.Metric != "concurrent" || reservation.Limit <= 0 || reservation.Amount <= 0 {
continue
}
key := fmt.Sprintf("%d:%s%d:%s", len(reservation.ScopeType), reservation.ScopeType, len(reservation.ScopeKey), reservation.ScopeKey)
if _, exists := lockKeySet[key]; exists {
continue
}
lockKeySet[key] = struct{}{}
lockKeys = append(lockKeys, key)
}
sort.Strings(lockKeys)
for _, key := range lockKeys {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, key); err != nil {
return RateLimitResult{}, err
}
}
result := RateLimitResult{}
for _, reservation := range reservations {
if reservation.Limit <= 0 || reservation.Amount <= 0 {
@@ -49,11 +72,11 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
reservation.WindowSeconds = 60
}
if reservation.Metric == "concurrent" {
leaseID, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation)
lease, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation)
if err != nil {
return RateLimitResult{}, err
}
result.LeaseIDs = append(result.LeaseIDs, leaseID)
result.Leases = append(result.Leases, lease)
continue
}
normalized, err := reserveCounterWindow(ctx, tx, taskID, attemptID, reservation)
@@ -65,7 +88,7 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
return result, tx.Commit(ctx)
}
func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (string, error) {
func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (ConcurrencyLease, error) {
if reservation.LeaseTTLSeconds <= 0 {
reservation.LeaseTTLSeconds = 120
}
@@ -83,10 +106,10 @@ WHERE scope_type = $1
reservation.ScopeKey,
reservation.LeaseTTLSeconds,
).Scan(&active, &nextAvailableAt); err != nil {
return "", err
return ConcurrencyLease{}, err
}
if active+reservation.Amount > reservation.Limit {
return "", &RateLimitExceededError{
return ConcurrencyLease{}, &RateLimitExceededError{
ScopeType: reservation.ScopeType,
ScopeKey: reservation.ScopeKey,
ScopeName: reservation.ScopeName,
@@ -117,9 +140,9 @@ RETURNING id::text`,
reservation.Amount,
reservation.LeaseTTLSeconds,
).Scan(&leaseID); err != nil {
return "", err
return ConcurrencyLease{}, err
}
return leaseID, nil
return ConcurrencyLease{ID: leaseID, TTL: time.Duration(reservation.LeaseTTLSeconds) * time.Second}, nil
}
func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (RateLimitReservation, error) {
@@ -232,13 +255,16 @@ func retryAfterUntil(when time.Time) time.Duration {
func concurrencyRetryAfter(leaseExpiresAt time.Time) time.Duration {
if leaseExpiresAt.IsZero() {
return time.Second
return 5 * time.Second
}
duration := time.Until(leaseExpiresAt)
if duration <= time.Second {
return time.Second
}
return time.Second
if duration > 5*time.Second {
return 5 * time.Second
}
return duration
}
func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error {
@@ -249,26 +275,57 @@ func (s *Store) ReleaseRateLimitReservations(ctx context.Context, reservations [
return s.finishRateLimitReservations(ctx, reservations, nil, "released", reason)
}
func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leaseIDs []string) error {
func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
leaseIDs := concurrencyLeaseIDs(leases)
if len(leaseIDs) == 0 {
return nil
}
for _, leaseID := range leaseIDs {
if leaseID == "" {
continue
}
if _, err := s.pool.Exec(ctx, `
_, err := s.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE id = $1::uuid AND released_at IS NULL`, leaseID); err != nil && !errors.Is(err, ErrRateLimited) {
return err
WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs)
return err
}
func (s *Store) RenewConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
leaseIDs := make([]string, 0, len(leases))
ttlSeconds := make([]int32, 0, len(leases))
for _, lease := range leases {
if lease.ID == "" {
continue
}
ttl := lease.TTL
if ttl <= 0 {
ttl = 120 * time.Second
}
seconds := int32(ttl / time.Second)
if seconds < 1 {
seconds = 1
}
leaseIDs = append(leaseIDs, lease.ID)
ttlSeconds = append(ttlSeconds, seconds)
}
if len(leaseIDs) == 0 {
return nil
}
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases lease
SET expires_at = now() + (renewal.ttl_seconds * interval '1 second')
FROM unnest($1::uuid[], $2::int[]) AS renewal(id, ttl_seconds)
WHERE lease.id = renewal.id
AND lease.released_at IS NULL
AND lease.expires_at > now()`, leaseIDs, ttlSeconds)
if err != nil {
return err
}
if tag.RowsAffected() != int64(len(leaseIDs)) {
return fmt.Errorf("%w: renewed %d of %d leases", ErrConcurrencyLeaseLost, tag.RowsAffected(), len(leaseIDs))
}
return nil
}
func (s *Store) AttachRateLimitResultToAttempt(ctx context.Context, attemptID string, result RateLimitResult) error {
if attemptID == "" || (len(result.Reservations) == 0 && len(result.LeaseIDs) == 0) {
if attemptID == "" || (len(result.Reservations) == 0 && len(result.Leases) == 0) {
return nil
}
tx, err := s.pool.Begin(ctx)
@@ -289,7 +346,8 @@ WHERE id = $1::uuid`, reservation.ReservationID, attemptID); err != nil {
return err
}
}
for _, leaseID := range result.LeaseIDs {
for _, lease := range result.Leases {
leaseID := lease.ID
if leaseID == "" {
continue
}
@@ -303,6 +361,16 @@ WHERE id = $1::uuid`, leaseID, attemptID); err != nil {
return tx.Commit(ctx)
}
func concurrencyLeaseIDs(leases []ConcurrencyLease) []string {
ids := make([]string, 0, len(leases))
for _, lease := range leases {
if lease.ID != "" {
ids = append(ids, lease.ID)
}
}
return ids
}
func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeRecoveryResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
@@ -0,0 +1,198 @@
package store
import (
"context"
"errors"
"os"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestConcurrencyLeaseReservationIsAtomicAcrossPools(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run concurrency lease PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
defer cancel()
first, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect first store: %v", err)
}
defer first.Close()
second, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect second store: %v", err)
}
defer second.Close()
scopeKey := "atomic-" + time.Now().UTC().Format("20060102150405.000000000")
taskIDs := createLeaseTestTasks(t, ctx, first, 256, scopeKey)
defer deleteLeaseTestTasks(t, first, taskIDs)
var successes atomic.Int64
var peak atomic.Int64
monitorCtx, stopMonitor := context.WithCancel(ctx)
var monitorWG sync.WaitGroup
monitorWG.Add(1)
go func() {
defer monitorWG.Done()
ticker := time.NewTicker(2 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-monitorCtx.Done():
return
case <-ticker.C:
var active int64
if err := first.Pool().QueryRow(monitorCtx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, scopeKey).Scan(&active); err == nil {
for active > peak.Load() && !peak.CompareAndSwap(peak.Load(), active) {
}
}
}
}
}()
var wg sync.WaitGroup
errs := make(chan error, len(taskIDs))
for index, taskID := range taskIDs {
wg.Add(1)
go func(index int, taskID string) {
defer wg.Done()
target := first
if index%2 == 1 {
target = second
}
_, err := target.ReserveRateLimits(ctx, taskID, "", []RateLimitReservation{{
ScopeType: "platform_model",
ScopeKey: scopeKey,
Metric: "concurrent",
Limit: 64,
Amount: 1,
LeaseTTLSeconds: 30,
}})
if err == nil {
successes.Add(1)
return
}
if !errors.Is(err, ErrRateLimited) {
errs <- err
}
}(index, taskID)
}
wg.Wait()
stopMonitor()
monitorWG.Wait()
close(errs)
for err := range errs {
t.Fatalf("unexpected reservation error: %v", err)
}
var active int64
if err := first.Pool().QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, scopeKey).Scan(&active); err != nil {
t.Fatalf("count active leases: %v", err)
}
if successes.Load() != 64 || active != 64 {
t.Fatalf("successful reservations=%d active leases=%d, want exactly 64", successes.Load(), active)
}
if peak.Load() > 64 {
t.Fatalf("active lease peak=%d, want <=64", peak.Load())
}
}
func TestConcurrencyLeaseRenewalExtendsAndReleases(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run concurrency lease PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
scopeKey := "renew-" + time.Now().UTC().Format("20060102150405.000000000")
taskIDs := createLeaseTestTasks(t, ctx, db, 1, scopeKey)
defer deleteLeaseTestTasks(t, db, taskIDs)
result, err := db.ReserveRateLimits(ctx, taskIDs[0], "", []RateLimitReservation{{
ScopeType: "platform_model",
ScopeKey: scopeKey,
Metric: "concurrent",
Limit: 1,
Amount: 1,
LeaseTTLSeconds: 2,
}})
if err != nil {
t.Fatalf("reserve short lease: %v", err)
}
time.Sleep(time.Second)
if err := db.RenewConcurrencyLeases(ctx, result.Leases); err != nil {
t.Fatalf("renew short lease: %v", err)
}
time.Sleep(1500 * time.Millisecond)
var active bool
if err := db.Pool().QueryRow(ctx, `
SELECT EXISTS (
SELECT 1 FROM gateway_concurrency_leases
WHERE id = $1::uuid AND released_at IS NULL AND expires_at > now()
)`, result.Leases[0].ID).Scan(&active); err != nil {
t.Fatalf("read renewed lease: %v", err)
}
if !active {
t.Fatal("renewed lease expired at its original TTL")
}
if err := db.ReleaseConcurrencyLeases(ctx, result.Leases); err != nil {
t.Fatalf("release renewed lease: %v", err)
}
}
func createLeaseTestTasks(t *testing.T, ctx context.Context, db *Store, count int, marker string) []string {
t.Helper()
rows, err := db.Pool().Query(ctx, `
INSERT INTO gateway_tasks (kind, run_mode, user_id, model, model_type, request, status, queue_key)
SELECT 'lease-test', 'simulation', $2, 'lease-test', 'text_generate', '{}'::jsonb, 'queued', $2
FROM generate_series(1, $1)
RETURNING id::text`, count, marker)
if err != nil {
t.Fatalf("create lease test tasks: %v", err)
}
defer rows.Close()
ids := make([]string, 0, count)
for rows.Next() {
var id string
if err := rows.Scan(&id); err != nil {
t.Fatalf("scan lease test task: %v", err)
}
ids = append(ids, id)
}
if err := rows.Err(); err != nil {
t.Fatalf("create lease test tasks: %v", err)
}
return ids
}
func deleteLeaseTestTasks(t *testing.T, db *Store, taskIDs []string) {
t.Helper()
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if _, err := db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, taskIDs); err != nil {
t.Errorf("delete lease test tasks: %v", err)
}
}
+11 -1
View File
@@ -111,6 +111,7 @@ type CreatePlatformModelInput struct {
PermissionConfig map[string]any `json:"permissionConfig"`
RetryPolicy map[string]any `json:"retryPolicy"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy"`
RateLimitPolicyMode string `json:"rateLimitPolicyMode" enums:"inherit,override"`
RuntimePolicySetID string `json:"runtimePolicySetId"`
RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride"`
Enabled bool `json:"enabled"`
@@ -130,6 +131,7 @@ type RuntimeModelCandidate struct {
DefaultDiscountFactor float64
PlatformRetryPolicy map[string]any
PlatformRateLimitPolicy map[string]any
BaseRateLimitPolicy map[string]any
PlatformPriority int
PlatformModelID string
BaseModelID string
@@ -153,8 +155,11 @@ type RuntimeModelCandidate struct {
ModelPricingRuleSetID string
ModelRetryPolicy map[string]any
ModelRateLimitPolicy map[string]any
ModelRateLimitPolicyMode string
RuntimePolicySetID string
RuntimePolicyExplicit bool
RuntimePolicyOverride map[string]any
RateLimitRuntimeOverride map[string]any
RuntimeRetryPolicy map[string]any
RuntimeRateLimitPolicy map[string]any
AutoDisablePolicy map[string]any
@@ -219,10 +224,15 @@ type RateLimitReservation struct {
}
type RateLimitResult struct {
LeaseIDs []string
Leases []ConcurrencyLease
Reservations []RateLimitReservation
}
type ConcurrencyLease struct {
ID string
TTL time.Duration
}
type CreateTaskAttemptInput struct {
TaskID string
AttemptNo int