jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 17:07:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: proactive rate limiting (RPM) + 429 short cooldown for sources

Browse Source 
- config.go: Source add RPM field (requests-per-minute cap, 0=unlimited)
- provider.go: RecordRateLimit() — 429 uses fixed 30s cooldown, not exponential
- provider.go: Throttle() — token-bucket proactive rate limiter, spaces requests
  at 60s/RPM interval, respects context cancellation
- provider.go: ReportStatus() — 429 -> RecordRateLimit, 5xx -> RecordFailure
- provider.go: Chat/ChatStream — wire Throttle after TryAcquire
- api.go: sourcePayload + RPM, buildSource passes RPM through
- ui/index.html: add RPM input field in source editor, bilingual i18n labels
- deploy.sh: backup old binary + rollback on healthcheck failure
- provider_test.go: TestModelStateRateLimitShortCooldown, TestThrottleSpacingAndCancel
- config.yaml: sensenova rpm: 12
This commit is contained in:
JianFeeeee
2026-08-24 02:00:25 +08:00
parent 3069cfce4e
commit 690f55c6f1
6 changed files with 429 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
internal/provider/provider.go
View File

@ -47,6 +47,13 @@ const (
backoffCapN = 10
)
// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
// rejection means "too fast", not "broken": unlike transport/5xx failures it
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
// locked out for up to 30 minutes while its allowance is untouched. A short
// fixed window lets the slot re-enter the rotation quickly.
const rateLimitCooldown = 30 * time.Second
func clampPref(a *atomic.Int64, lo, hi int64) {
for {
cur := a.Load()
@ -100,6 +107,17 @@ func (s *ModelState) RecordFailure(auth bool) {
clampPref(&s.pref, prefMin, prefMax)
}
// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
// exponential schedule, so a merely-throttled source recovers in seconds and
// its remaining quota stays usable. The preference penalty still applies so
// other slots are preferred while cooling.
func (s *ModelState) RecordRateLimit() {
s.failCount.Add(1)
s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix())
s.pref.Add(-prefFailStep)
clampPref(&s.pref, prefMin, prefMax)
}
// RecordSuccess resets the failure counter and cooldown and bumps the
// preference score by one.
func (s *ModelState) RecordSuccess() {
@ -139,6 +157,12 @@ type Provider struct {
mu sync.Mutex
sem chan struct{}
states map[string]*ModelState // key = model id
// proactive rate limiting: nextOK is the earliest unix-nano time the next
// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
// cfg.RPM <= 0 (unlimited).
rateMu sync.Mutex
nextOK int64
rpmGap time.Duration
lastProbe struct {
ok bool
err string
@ -178,6 +202,9 @@ func New(cfg config.Source, vm *lua.VM) *Provider {
if cfg.MaxConcurrent <= 0 {
p.sem = nil
}
if cfg.RPM > 0 {
p.rpmGap = time.Minute / time.Duration(cfg.RPM)
}
for _, m := range cfg.Models {
p.states[m.ID] = &ModelState{}
}
@ -491,15 +518,20 @@ func (p *Provider) RecordSuccess(model string) {
// ReportStatus records an upstream HTTP status for the given model and drives
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
// with doubled penalty; 429 and 5xx → normal exponential backoff. Other codes
// (400 client schema errors, 402 billing errors) are not penalized here —
// they surface via the status page / audit instead.
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
// "broken", so the slot must re-enter rotation quickly instead of escalating
// to a 30-minute lockout that wastes remaining quota.
func (p *Provider) ReportStatus(model string, code int) {
if code == 401 || code == 403 {
p.RecordFailure(model, code)
return
}
if code >= 500 || code == 429 {
if code == 429 {
p.state(model).RecordRateLimit()
return
}
if code >= 500 {
p.RecordFailure(model, code)
}
}
@ -586,6 +618,37 @@ func (p *Provider) Release() {
<-p.sem
}
// ---- proactive rate limiting ----
// Throttle blocks until this source's rate-limit window allows the next
// request, spacing outbound requests at least rpmGap apart (60s/RPM). It is a
// no-op when no rpm cap is configured. The wait happens AFTER the concurrency
// slot is taken, so a queued request still counts against max_concurrent —
// bounded by the caller's ctx (queue timeout / client disconnect).
func (p *Provider) Throttle(ctx context.Context) error {
if p.rpmGap <= 0 {
return nil
}
p.rateMu.Lock()
now := time.Now().UnixNano()
wait := p.nextOK - now
if wait > 0 {
p.nextOK += int64(p.rpmGap) // reserve the next window for the follower
p.rateMu.Unlock()
t := time.NewTimer(time.Duration(wait))
defer t.Stop()
select {
case <-t.C:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
p.nextOK = now + int64(p.rpmGap)
p.rateMu.Unlock()
return nil
}
// ---- request construction ----
func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
@ -630,6 +693,9 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
return nil, err
}
defer p.Release()
if err := p.Throttle(ctx); err != nil {
return nil, err
}
model := p.ModelFor(req.Model)
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
@ -684,6 +750,10 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
if err := p.TryAcquire(ctx); err != nil {
return nil, err
}
if err := p.Throttle(ctx); err != nil {
p.Release()
return nil, err
}
model := p.ModelFor(req.Model)
req.Stream = true
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)

71
internal/provider/provider_test.go
View File

@ -245,6 +245,77 @@ func TestModelStateCooldownAndRecovery(t *testing.T) {
}
}
func TestModelStateRateLimitShortCooldown(t *testing.T) {
// 429 must NOT use the exponential backoff schedule: a rate-limited but
// quota-rich source has to re-enter rotation after the short fixed window.
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
st := p.state("m")
// repeated 429s must stay at the fixed short cooldown, never escalate
for i := 0; i < 15; i++ {
p.ReportStatus("m", 429)
}
until := st.CooldownUntil()
want := time.Now().Add(rateLimitCooldown).Unix()
if until < want-2 || until > want+2 {
t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown)
}
if st.FailCount() != 15 {
t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount())
}
if st.Pref() != -15*int64(prefFailStep) && st.Pref() > int64(prefMin) {
t.Fatalf("pref = %d", st.Pref())
}
if p.ModelAvailable("m") {
t.Fatal("model must be cooling right after a 429")
}
}
func TestThrottleSpacingAndCancel(t *testing.T) {
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
p.cfg.RPM = 120 // gap = 500ms
p.rpmGap = time.Minute / time.Duration(p.cfg.RPM)
start := time.Now()
for i := 0; i < 3; i++ {
if err := p.Throttle(context.Background()); err != nil {
t.Fatalf("throttle %d: %v", i, err)
}
}
elapsed := time.Since(start)
// first call passes immediately, the next two wait one gap each
want := 2 * time.Minute / time.Duration(p.cfg.RPM)
if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second {
t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want)
}
// unlimited source: Throttle is a no-op
p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m"))
if err := p2.Throttle(context.Background()); err != nil {
t.Fatalf("unlimited throttle: %v", err)
}
// cancelled context aborts a pending window reservation
p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m"))
p3.cfg.RPM = 6 // 10s gap
p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM)
// first call passes immediately (fresh provider) and reserves the next window
if err := p3.Throttle(context.Background()); err != nil {
t.Fatalf("first throttle: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
// second call must wait ~10s for the reserved window; the short deadline aborts it
start = time.Now()
err := p3.Throttle(ctx)
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err)
}
if time.Since(start) > time.Second {
t.Fatalf("cancel took %v, want fast abort", time.Since(start))
}
}
func TestTryAcquire(t *testing.T) {
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
p.cfg.MaxConcurrent = 1