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:
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user