97 lines
3.7 KiB
Go
97 lines
3.7 KiB
Go
package pipeline
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// ratelimit.go: the per-model rate-guard for the parallel waves (WS1 §1б, ревью-2 F4). The per-CALL
|
||
// 429/Retry-After backoff already lives in the transport adapter (httpllm.go); the gap the waves
|
||
// open is N-CONCURRENCY against a rate-limited model — mistral-large-latest fails ~48% of calls under
|
||
// parallelism (a token-bucket / concurrency cap, tier-dependent — 00-provider-quirks §Транспорт),
|
||
// while grok is 0%. A model caps its own wave concurrency (max_concurrency) and, optionally, paces
|
||
// call STARTS (min_interval). This is a TRANSPORT axis: it never touches the request bytes, so it is
|
||
// wire-neutral and NOT snapshot-folded. Guards are built once in W0, read-only in the waves; each
|
||
// guard's internals (a buffered-channel semaphore + a mutex-protected pacer) are concurrency-safe,
|
||
// so acquiring from N wave goroutines is race-clean.
|
||
|
||
// rateGuard limits concurrent calls to one model (+ optional pacing). The zero/nil guard is a no-op
|
||
// (unlimited), so an unconfigured model — or a nil map entry — costs nothing.
|
||
type rateGuard struct {
|
||
sem chan struct{} // buffered to max_concurrency; nil = unlimited
|
||
minInterval time.Duration // 0 = no pacing
|
||
mu sync.Mutex // guards last
|
||
last time.Time // start time of the most recent acquired call (pacer)
|
||
}
|
||
|
||
// newRateGuard builds a guard. maxConcurrency ≤ 0 → unlimited concurrency; minInterval ≤ 0 → no
|
||
// pacing. Both off → newRateGuard returns a guard that acquire treats as a pure no-op.
|
||
func newRateGuard(maxConcurrency int, minInterval time.Duration) *rateGuard {
|
||
g := &rateGuard{minInterval: minInterval}
|
||
if maxConcurrency > 0 {
|
||
g.sem = make(chan struct{}, maxConcurrency)
|
||
}
|
||
return g
|
||
}
|
||
|
||
// acquire blocks until a concurrency slot is free and, if pacing is on, until the min-interval since
|
||
// the last acquired call has elapsed — both respecting ctx cancellation. It returns a release func
|
||
// the caller MUST defer (a no-op when acquire failed or the guard is unlimited). A nil guard (no
|
||
// config) acquires instantly. The pacer holds the mutex only to read/stamp `last`; the wait itself
|
||
// is outside the lock, so a paced model still admits one caller per interval without serializing the
|
||
// lock for the whole wait.
|
||
func (g *rateGuard) acquire(ctx context.Context) (func(), error) {
|
||
if g == nil {
|
||
return func() {}, nil
|
||
}
|
||
if g.sem != nil {
|
||
select {
|
||
case g.sem <- struct{}{}:
|
||
case <-ctx.Done():
|
||
return func() {}, ctx.Err()
|
||
}
|
||
}
|
||
release := func() {
|
||
if g.sem != nil {
|
||
<-g.sem
|
||
}
|
||
}
|
||
if g.minInterval > 0 {
|
||
for {
|
||
g.mu.Lock()
|
||
wait := time.Until(g.last.Add(g.minInterval))
|
||
if wait <= 0 {
|
||
g.last = time.Now()
|
||
g.mu.Unlock()
|
||
break
|
||
}
|
||
g.mu.Unlock()
|
||
select {
|
||
case <-time.After(wait):
|
||
// re-check: another caller may have taken the slot while we slept.
|
||
case <-ctx.Done():
|
||
release()
|
||
return func() {}, ctx.Err()
|
||
}
|
||
}
|
||
}
|
||
return release, nil
|
||
}
|
||
|
||
// buildRateGuards constructs a guard for every reachable model in W0 (read-only in the waves).
|
||
// A model with no rate_limit config gets an unlimited (no-op) guard, so rateGuard(model) is always
|
||
// non-nil and lookup-safe.
|
||
func (r *Runner) buildRateGuards() {
|
||
r.rateGuards = make(map[string]*rateGuard, len(r.Models.Models))
|
||
for _, m := range r.reachableModels() {
|
||
rl := r.Models.Models[m].RateLimit
|
||
r.rateGuards[m] = newRateGuard(rl.MaxConcurrency, time.Duration(rl.MinIntervalMS)*time.Millisecond)
|
||
}
|
||
}
|
||
|
||
// rateGuard returns the pre-built guard for a model (nil-safe: acquire treats a missing/nil guard as
|
||
// unlimited, so a model reached before buildRateGuards ran simply is not throttled).
|
||
func (r *Runner) rateGuard(model string) *rateGuard {
|
||
return r.rateGuards[model]
|
||
}
|