feat: add RPM getter methods and schedulability check to Account model
This commit is contained in:
@@ -1137,6 +1137,77 @@ func (a *Account) GetSessionIdleTimeoutMinutes() int {
|
|||||||
return 5
|
return 5
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// GetBaseRPM 获取基础 RPM 限制
|
||||||
|
// 返回 0 表示未启用
|
||||||
|
func (a *Account) GetBaseRPM() int {
|
||||||
|
if a.Extra == nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
if v, ok := a.Extra["base_rpm"]; ok {
|
||||||
|
return parseExtraInt(v)
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetRPMStrategy 获取 RPM 策略
|
||||||
|
// "tiered" = 三区模型(默认), "sticky_exempt" = 粘性豁免
|
||||||
|
func (a *Account) GetRPMStrategy() string {
|
||||||
|
if a.Extra == nil {
|
||||||
|
return "tiered"
|
||||||
|
}
|
||||||
|
if v, ok := a.Extra["rpm_strategy"]; ok {
|
||||||
|
if s, ok := v.(string); ok && s == "sticky_exempt" {
|
||||||
|
return "sticky_exempt"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return "tiered"
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetRPMStickyBuffer 获取 RPM 粘性缓冲数量
|
||||||
|
// tiered 模式下的黄区大小,默认为 base_rpm 的 20%(至少 1)
|
||||||
|
func (a *Account) GetRPMStickyBuffer() int {
|
||||||
|
if a.Extra == nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
if v, ok := a.Extra["rpm_sticky_buffer"]; ok {
|
||||||
|
val := parseExtraInt(v)
|
||||||
|
if val > 0 {
|
||||||
|
return val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
base := a.GetBaseRPM()
|
||||||
|
buffer := base / 5
|
||||||
|
if buffer < 1 && base > 0 {
|
||||||
|
buffer = 1
|
||||||
|
}
|
||||||
|
return buffer
|
||||||
|
}
|
||||||
|
|
||||||
|
// CheckRPMSchedulability 根据当前 RPM 计数检查调度状态
|
||||||
|
// 复用 WindowCostSchedulability 三态:Schedulable / StickyOnly / NotSchedulable
|
||||||
|
func (a *Account) CheckRPMSchedulability(currentRPM int) WindowCostSchedulability {
|
||||||
|
baseRPM := a.GetBaseRPM()
|
||||||
|
if baseRPM <= 0 {
|
||||||
|
return WindowCostSchedulable
|
||||||
|
}
|
||||||
|
|
||||||
|
if currentRPM < baseRPM {
|
||||||
|
return WindowCostSchedulable
|
||||||
|
}
|
||||||
|
|
||||||
|
strategy := a.GetRPMStrategy()
|
||||||
|
if strategy == "sticky_exempt" {
|
||||||
|
return WindowCostStickyOnly // 粘性豁免无红区
|
||||||
|
}
|
||||||
|
|
||||||
|
// tiered: 黄区 + 红区
|
||||||
|
buffer := a.GetRPMStickyBuffer()
|
||||||
|
if currentRPM < baseRPM+buffer {
|
||||||
|
return WindowCostStickyOnly
|
||||||
|
}
|
||||||
|
return WindowCostNotSchedulable
|
||||||
|
}
|
||||||
|
|
||||||
// CheckWindowCostSchedulability 根据当前窗口费用检查调度状态
|
// CheckWindowCostSchedulability 根据当前窗口费用检查调度状态
|
||||||
// - 费用 < 阈值: WindowCostSchedulable(可正常调度)
|
// - 费用 < 阈值: WindowCostSchedulable(可正常调度)
|
||||||
// - 费用 >= 阈值 且 < 阈值+预留: WindowCostStickyOnly(仅粘性会话)
|
// - 费用 >= 阈值 且 < 阈值+预留: WindowCostStickyOnly(仅粘性会话)
|
||||||
|
|||||||
73
backend/internal/service/account_rpm_test.go
Normal file
73
backend/internal/service/account_rpm_test.go
Normal file
@@ -0,0 +1,73 @@
|
|||||||
|
package service
|
||||||
|
|
||||||
|
import "testing"
|
||||||
|
|
||||||
|
func TestGetBaseRPM(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
extra map[string]any
|
||||||
|
expected int
|
||||||
|
}{
|
||||||
|
{"nil extra", nil, 0},
|
||||||
|
{"no key", map[string]any{}, 0},
|
||||||
|
{"zero", map[string]any{"base_rpm": 0}, 0},
|
||||||
|
{"int value", map[string]any{"base_rpm": 15}, 15},
|
||||||
|
{"float value", map[string]any{"base_rpm": 15.0}, 15},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
a := &Account{Extra: tt.extra}
|
||||||
|
if got := a.GetBaseRPM(); got != tt.expected {
|
||||||
|
t.Errorf("GetBaseRPM() = %d, want %d", got, tt.expected)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetRPMStrategy(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
extra map[string]any
|
||||||
|
expected string
|
||||||
|
}{
|
||||||
|
{"nil extra", nil, "tiered"},
|
||||||
|
{"no key", map[string]any{}, "tiered"},
|
||||||
|
{"tiered", map[string]any{"rpm_strategy": "tiered"}, "tiered"},
|
||||||
|
{"sticky_exempt", map[string]any{"rpm_strategy": "sticky_exempt"}, "sticky_exempt"},
|
||||||
|
{"invalid", map[string]any{"rpm_strategy": "foobar"}, "tiered"},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
a := &Account{Extra: tt.extra}
|
||||||
|
if got := a.GetRPMStrategy(); got != tt.expected {
|
||||||
|
t.Errorf("GetRPMStrategy() = %q, want %q", got, tt.expected)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCheckRPMSchedulability(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
extra map[string]any
|
||||||
|
currentRPM int
|
||||||
|
expected WindowCostSchedulability
|
||||||
|
}{
|
||||||
|
{"disabled", map[string]any{}, 100, WindowCostSchedulable},
|
||||||
|
{"green zone", map[string]any{"base_rpm": 15}, 10, WindowCostSchedulable},
|
||||||
|
{"yellow zone tiered", map[string]any{"base_rpm": 15}, 15, WindowCostStickyOnly},
|
||||||
|
{"red zone tiered", map[string]any{"base_rpm": 15}, 18, WindowCostNotSchedulable},
|
||||||
|
{"sticky_exempt at limit", map[string]any{"base_rpm": 15, "rpm_strategy": "sticky_exempt"}, 15, WindowCostStickyOnly},
|
||||||
|
{"sticky_exempt over limit", map[string]any{"base_rpm": 15, "rpm_strategy": "sticky_exempt"}, 100, WindowCostStickyOnly},
|
||||||
|
{"custom buffer", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 14, WindowCostStickyOnly},
|
||||||
|
{"custom buffer red", map[string]any{"base_rpm": 10, "rpm_sticky_buffer": 5}, 15, WindowCostNotSchedulable},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
a := &Account{Extra: tt.extra}
|
||||||
|
if got := a.CheckRPMSchedulability(tt.currentRPM); got != tt.expected {
|
||||||
|
t.Errorf("CheckRPMSchedulability(%d) = %d, want %d", tt.currentRPM, got, tt.expected)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user