core/go-ratelimit
8
0
Fork 0
Code Issues 4 Pull requests Projects Releases Packages Wiki Activity Actions

refactor: apply go fix modernizers for Go 1.26

Browse source 
Automated fixes: interface{} → any, range-over-int, t.Context(),
wg.Go(), strings.SplitSeq, strings.Builder, slices.Contains,
maps helpers, min/max builtins.

Co-Authored-By: Virgil <virgil@lethean.io>
This commit is contained in:
Snider 2026-02-22 21:00:17 +00:00
parent 0afcf0fb64
commit 2246fea10f
3 changed files with 72 additions and 105 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
ratelimit.go
View file

@ -6,6 +6,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"maps"
"net/http" "net/http"
"os" "os"
"path/filepath" "path/filepath"
@ -38,8 +39,8 @@ type ModelQuota struct {
// ProviderProfile bundles model quotas for a provider. // ProviderProfile bundles model quotas for a provider.
type ProviderProfile struct { type ProviderProfile struct {
Provider Provider `yaml:"provider"` Provider Provider `yaml:"provider"`
Models map[string]ModelQuota `yaml:"models"` Models map[string]ModelQuota `yaml:"models"`
} }
// Config controls RateLimiter initialisation. // Config controls RateLimiter initialisation.
@ -101,12 +102,12 @@ func DefaultProfiles() map[Provider]ProviderProfile {
ProviderOpenAI: { ProviderOpenAI: {
Provider: ProviderOpenAI, Provider: ProviderOpenAI,
Models: map[string]ModelQuota{ Models: map[string]ModelQuota{
"gpt-4o": {MaxRPM: 500, MaxTPM: 30000, MaxRPD: 0}, "gpt-4o": {MaxRPM: 500, MaxTPM: 30000, MaxRPD: 0},
"gpt-4o-mini": {MaxRPM: 500, MaxTPM: 200000, MaxRPD: 0}, "gpt-4o-mini": {MaxRPM: 500, MaxTPM: 200000, MaxRPD: 0},
"gpt-4-turbo": {MaxRPM: 500, MaxTPM: 30000, MaxRPD: 0}, "gpt-4-turbo": {MaxRPM: 500, MaxTPM: 30000, MaxRPD: 0},
"o1": {MaxRPM: 500, MaxTPM: 30000, MaxRPD: 0}, "o1": {MaxRPM: 500, MaxTPM: 30000, MaxRPD: 0},
"o1-mini": {MaxRPM: 500, MaxTPM: 200000, MaxRPD: 0}, "o1-mini": {MaxRPM: 500, MaxTPM: 200000, MaxRPD: 0},
"o3-mini": {MaxRPM: 500, MaxTPM: 200000, MaxRPD: 0}, "o3-mini": {MaxRPM: 500, MaxTPM: 200000, MaxRPD: 0},
}, },
}, },
ProviderAnthropic: { ProviderAnthropic: {
@ -119,7 +120,7 @@ func DefaultProfiles() map[Provider]ProviderProfile {
}, },
ProviderLocal: { ProviderLocal: {
Provider: ProviderLocal, Provider: ProviderLocal,
Models: map[string]ModelQuota{ Models: map[string]ModelQuota{
// Local inference has no external rate limits by default. // Local inference has no external rate limits by default.
// Users can override per-model if their hardware requires throttling. // Users can override per-model if their hardware requires throttling.
}, },
@ -164,16 +165,12 @@ func NewWithConfig(cfg Config) (*RateLimiter, error) {
for _, p := range providers { for _, p := range providers {
if profile, ok := profiles[p]; ok { if profile, ok := profiles[p]; ok {
for model, quota := range profile.Models { maps.Copy(rl.Quotas, profile.Models)
rl.Quotas[model] = quota
}
} }
} }
// Merge explicit quotas on top (allows overrides) // Merge explicit quotas on top (allows overrides)
for model, quota := range cfg.Quotas { maps.Copy(rl.Quotas, cfg.Quotas)
rl.Quotas[model] = quota
}
return rl, nil return rl, nil
} }
@ -193,9 +190,7 @@ func (rl *RateLimiter) AddProvider(provider Provider) {
profiles := DefaultProfiles() profiles := DefaultProfiles()
if profile, ok := profiles[provider]; ok { if profile, ok := profiles[provider]; ok {
for model, quota := range profile.Models { maps.Copy(rl.Quotas, profile.Models)
rl.Quotas[model] = quota
}
} }
} }
@ -227,18 +222,14 @@ func (rl *RateLimiter) loadSQLite() error {
return err return err
} }
// Merge loaded quotas (loaded quotas override in-memory defaults). // Merge loaded quotas (loaded quotas override in-memory defaults).
for model, q := range quotas { maps.Copy(rl.Quotas, quotas)
rl.Quotas[model] = q
}
state, err := rl.sqlite.loadState() state, err := rl.sqlite.loadState()
if err != nil { if err != nil {
return err return err
} }
// Replace in-memory state with persisted state. // Replace in-memory state with persisted state.
for model, s := range state { maps.Copy(rl.State, state)
rl.State[model] = s
}
return nil return nil
} }
@ -545,14 +536,10 @@ func NewWithSQLiteConfig(dbPath string, cfg Config) (*RateLimiter, error) {
} }
for _, p := range providers { for _, p := range providers {
if profile, ok := profiles[p]; ok { if profile, ok := profiles[p]; ok {
for model, quota := range profile.Models { maps.Copy(rl.Quotas, profile.Models)
rl.Quotas[model] = quota
}
} }
} }
for model, quota := range cfg.Quotas { maps.Copy(rl.Quotas, cfg.Quotas)
rl.Quotas[model] = quota
}
return rl, nil return rl, nil
} }

114
ratelimit_test.go
View file

@ -46,7 +46,7 @@ func TestCanSend(t *testing.T) {
model := "test-unlimited" model := "test-unlimited"
rl.Quotas[model] = ModelQuota{MaxRPM: 0, MaxTPM: 0, MaxRPD: 0} rl.Quotas[model] = ModelQuota{MaxRPM: 0, MaxTPM: 0, MaxRPD: 0}
for i := 0; i < 1000; i++ { for range 1000 {
rl.RecordUsage(model, 100, 100) rl.RecordUsage(model, 100, 100)
} }
assert.True(t, rl.CanSend(model, 999999), "unlimited model should always allow sends") assert.True(t, rl.CanSend(model, 999999), "unlimited model should always allow sends")
@ -122,7 +122,7 @@ func TestCanSend(t *testing.T) {
model := "test-rpd-unlimited" model := "test-rpd-unlimited"
rl.Quotas[model] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 0} rl.Quotas[model] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 0}
for i := 0; i < 100; i++ { for range 100 {
rl.RecordUsage(model, 1, 1) rl.RecordUsage(model, 1, 1)
} }
assert.True(t, rl.CanSend(model, 1), "RPD=0 should mean unlimited daily requests") assert.True(t, rl.CanSend(model, 1), "RPD=0 should mean unlimited daily requests")
@ -655,16 +655,14 @@ func TestConcurrentAccess(t *testing.T) {
goroutines := 20 goroutines := 20
opsPerGoroutine := 50 opsPerGoroutine := 50
for i := 0; i < goroutines; i++ { for range goroutines {
wg.Add(1) wg.Go(func() {
go func() { for range opsPerGoroutine {
defer wg.Done()
for j := 0; j < opsPerGoroutine; j++ {
rl.CanSend(model, 10) rl.CanSend(model, 10)
rl.RecordUsage(model, 5, 5) rl.RecordUsage(model, 5, 5)
rl.Stats(model) rl.Stats(model)
} }
}() })
} }
wg.Wait() wg.Wait()
@ -682,36 +680,30 @@ func TestConcurrentResetAndRecord(t *testing.T) {
var wg sync.WaitGroup var wg sync.WaitGroup
// Writers // Writers
for i := 0; i < 5; i++ { for range 5 {
wg.Add(1) wg.Go(func() {
go func() { for range 100 {
defer wg.Done()
for j := 0; j < 100; j++ {
rl.RecordUsage(model, 1, 1) rl.RecordUsage(model, 1, 1)
} }
}() })
} }
// Resetters // Resetters
for i := 0; i < 3; i++ { for range 3 {
wg.Add(1) wg.Go(func() {
go func() { for range 20 {
defer wg.Done()
for j := 0; j < 20; j++ {
rl.Reset(model) rl.Reset(model)
} }
}() })
} }
// Readers // Readers
for i := 0; i < 5; i++ { for range 5 {
wg.Add(1) wg.Go(func() {
go func() { for range 100 {
defer wg.Done()
for j := 0; j < 100; j++ {
rl.AllStats() rl.AllStats()
} }
}() })
} }
wg.Wait() wg.Wait()
@ -763,7 +755,7 @@ func BenchmarkCanSend(b *testing.B) {
now := time.Now() now := time.Now()
entries := make([]time.Time, 1000) entries := make([]time.Time, 1000)
tokens := make([]TokenEntry, 1000) tokens := make([]TokenEntry, 1000)
for i := 0; i < 1000; i++ { for i := range 1000 {
t := now.Add(-time.Duration(i) * time.Millisecond * 50) // spread over ~50 seconds t := now.Add(-time.Duration(i) * time.Millisecond * 50) // spread over ~50 seconds
entries[i] = t entries[i] = t
tokens[i] = TokenEntry{Time: t, Count: 10} tokens[i] = TokenEntry{Time: t, Count: 10}
@ -799,7 +791,7 @@ func BenchmarkCanSendConcurrent(b *testing.B) {
// Populate window // Populate window
now := time.Now() now := time.Now()
for i := 0; i < 100; i++ { for range 100 {
rl.State[model] = &UsageStats{DayStart: now} rl.State[model] = &UsageStats{DayStart: now}
rl.RecordUsage(model, 10, 10) rl.RecordUsage(model, 10, 10)
} }
@ -1014,7 +1006,7 @@ func TestSetQuota(t *testing.T) {
rl := newTestLimiter(t) rl := newTestLimiter(t)
var wg sync.WaitGroup var wg sync.WaitGroup
for i := 0; i < 10; i++ { for i := range 10 {
wg.Add(1) wg.Add(1)
go func(n int) { go func(n int) {
defer wg.Done() defer wg.Done()
@ -1080,7 +1072,7 @@ func TestAddProvider(t *testing.T) {
wg.Add(1) wg.Add(1)
go func(prov Provider) { go func(prov Provider) {
defer wg.Done() defer wg.Done()
for i := 0; i < 10; i++ { for range 10 {
rl.AddProvider(prov) rl.AddProvider(prov)
} }
}(p) }(p)
@ -1114,7 +1106,7 @@ func TestConcurrentMultipleModels(t *testing.T) {
wg.Add(1) wg.Add(1)
go func(model string) { go func(model string) {
defer wg.Done() defer wg.Done()
for i := 0; i < iterations; i++ { for range iterations {
rl.CanSend(model, 10) rl.CanSend(model, 10)
rl.RecordUsage(model, 10, 10) rl.RecordUsage(model, 10, 10)
rl.Stats(model) rl.Stats(model)
@ -1142,26 +1134,22 @@ func TestConcurrentPersistAndLoad(t *testing.T) {
var wg sync.WaitGroup var wg sync.WaitGroup
// Writers + persist // Writers + persist
for g := 0; g < 3; g++ { for range 3 {
wg.Add(1) wg.Go(func() {
go func() { for range 50 {
defer wg.Done()
for i := 0; i < 50; i++ {
rl.RecordUsage(model, 10, 10) rl.RecordUsage(model, 10, 10)
_ = rl.Persist() _ = rl.Persist()
} }
}() })
} }
// Loaders // Loaders
for g := 0; g < 3; g++ { for range 3 {
wg.Add(1) wg.Go(func() {
go func() { for range 50 {
defer wg.Done()
for i := 0; i < 50; i++ {
_ = rl.Load() _ = rl.Load()
} }
}() })
} }
wg.Wait() wg.Wait()
@ -1181,21 +1169,19 @@ func TestConcurrentAllStatsAndRecordUsage(t *testing.T) {
wg.Add(1) wg.Add(1)
go func(model string) { go func(model string) {
defer wg.Done() defer wg.Done()
for i := 0; i < 100; i++ { for range 100 {
rl.RecordUsage(model, 10, 10) rl.RecordUsage(model, 10, 10)
} }
}(m) }(m)
} }
// Read AllStats concurrently // Read AllStats concurrently
for g := 0; g < 3; g++ { for range 3 {
wg.Add(1) wg.Go(func() {
go func() { for range 50 {
defer wg.Done()
for i := 0; i < 50; i++ {
_ = rl.AllStats() _ = rl.AllStats()
} }
}() })
} }
wg.Wait() wg.Wait()
@ -1208,25 +1194,21 @@ func TestConcurrentWaitForCapacityAndRecordUsage(t *testing.T) {
var wg sync.WaitGroup var wg sync.WaitGroup
for g := 0; g < 5; g++ { for range 5 {
wg.Add(1) wg.Go(func() {
go func() {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel() defer cancel()
_ = rl.WaitForCapacity(ctx, model, 10) _ = rl.WaitForCapacity(ctx, model, 10)
}() })
} }
// Record usage concurrently // Record usage concurrently
for g := 0; g < 5; g++ { for range 5 {
wg.Add(1) wg.Go(func() {
go func() { for range 20 {
defer wg.Done()
for i := 0; i < 20; i++ {
rl.RecordUsage(model, 10, 10) rl.RecordUsage(model, 10, 10)
} }
}() })
} }
wg.Wait() wg.Wait()
@ -1242,12 +1224,12 @@ func BenchmarkCanSendWithPrune(b *testing.B) {
// Pre-fill with a mix of old and new entries to trigger pruning // Pre-fill with a mix of old and new entries to trigger pruning
now := time.Now() now := time.Now()
rl.State[model] = &UsageStats{DayStart: now} rl.State[model] = &UsageStats{DayStart: now}
for i := 0; i < 500; i++ { for range 500 {
old := now.Add(-2 * time.Minute) old := now.Add(-2 * time.Minute)
rl.State[model].Requests = append(rl.State[model].Requests, old) rl.State[model].Requests = append(rl.State[model].Requests, old)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: old, Count: 100}) rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: old, Count: 100})
} }
for i := 0; i < 500; i++ { for i := range 500 {
recent := now.Add(-time.Duration(i) * time.Millisecond * 100) recent := now.Add(-time.Duration(i) * time.Millisecond * 100)
rl.State[model].Requests = append(rl.State[model].Requests, recent) rl.State[model].Requests = append(rl.State[model].Requests, recent)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: recent, Count: 100}) rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: recent, Count: 100})
@ -1267,7 +1249,7 @@ func BenchmarkStats(b *testing.B) {
now := time.Now() now := time.Now()
rl.State[model] = &UsageStats{DayStart: now, DayCount: 500} rl.State[model] = &UsageStats{DayStart: now, DayCount: 500}
for i := 0; i < 1000; i++ { for i := range 1000 {
t := now.Add(-time.Duration(i) * time.Millisecond * 50) t := now.Add(-time.Duration(i) * time.Millisecond * 50)
rl.State[model].Requests = append(rl.State[model].Requests, t) rl.State[model].Requests = append(rl.State[model].Requests, t)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100}) rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100})
@ -1287,7 +1269,7 @@ func BenchmarkAllStats(b *testing.B) {
for _, m := range models { for _, m := range models {
rl.Quotas[m] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 100000} rl.Quotas[m] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 100000}
rl.State[m] = &UsageStats{DayStart: now, DayCount: 200} rl.State[m] = &UsageStats{DayStart: now, DayCount: 200}
for i := 0; i < 200; i++ { for i := range 200 {
t := now.Add(-time.Duration(i) * time.Millisecond * 250) t := now.Add(-time.Duration(i) * time.Millisecond * 250)
rl.State[m].Requests = append(rl.State[m].Requests, t) rl.State[m].Requests = append(rl.State[m].Requests, t)
rl.State[m].Tokens = append(rl.State[m].Tokens, TokenEntry{Time: t, Count: 100}) rl.State[m].Tokens = append(rl.State[m].Tokens, TokenEntry{Time: t, Count: 100})
@ -1311,7 +1293,7 @@ func BenchmarkPersist(b *testing.B) {
now := time.Now() now := time.Now()
rl.State[model] = &UsageStats{DayStart: now, DayCount: 100} rl.State[model] = &UsageStats{DayStart: now, DayCount: 100}
for i := 0; i < 100; i++ { for i := range 100 {
t := now.Add(-time.Duration(i) * time.Second) t := now.Add(-time.Duration(i) * time.Second)
rl.State[model].Requests = append(rl.State[model].Requests, t) rl.State[model].Requests = append(rl.State[model].Requests, t)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100}) rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100})

16
sqlite_test.go
View file

@ -311,7 +311,7 @@ func TestSQLiteRecordUsageThenPersistReload_Good(t *testing.T) {
rl.Quotas[model] = ModelQuota{MaxRPM: 100, MaxTPM: 100000, MaxRPD: 1000} rl.Quotas[model] = ModelQuota{MaxRPM: 100, MaxTPM: 100000, MaxRPD: 1000}
// Record multiple usages. // Record multiple usages.
for i := 0; i < 10; i++ { for range 10 {
rl.RecordUsage(model, 50, 50) rl.RecordUsage(model, 50, 50)
} }
@ -363,16 +363,14 @@ func TestSQLiteConcurrent_Good(t *testing.T) {
// Concurrent RecordUsage + CanSend + Persist (no Load, which would // Concurrent RecordUsage + CanSend + Persist (no Load, which would
// overwrite in-memory state and lose recordings between cycles). // overwrite in-memory state and lose recordings between cycles).
for i := 0; i < goroutines; i++ { for range goroutines {
wg.Add(1) wg.Go(func() {
go func() { for range opsPerGoroutine {
defer wg.Done()
for j := 0; j < opsPerGoroutine; j++ {
rl.RecordUsage(model, 5, 5) rl.RecordUsage(model, 5, 5)
rl.CanSend(model, 10) rl.CanSend(model, 10)
_ = rl.Persist() _ = rl.Persist()
} }
}() })
} }
wg.Wait() wg.Wait()
@ -667,7 +665,7 @@ func BenchmarkSQLitePersist(b *testing.B) {
now := time.Now() now := time.Now()
rl.State[model] = &UsageStats{DayStart: now, DayCount: 100} rl.State[model] = &UsageStats{DayStart: now, DayCount: 100}
for i := 0; i < 100; i++ { for i := range 100 {
t := now.Add(-time.Duration(i) * time.Second) t := now.Add(-time.Duration(i) * time.Second)
rl.State[model].Requests = append(rl.State[model].Requests, t) rl.State[model].Requests = append(rl.State[model].Requests, t)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100}) rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100})
@ -692,7 +690,7 @@ func BenchmarkSQLiteLoad(b *testing.B) {
now := time.Now() now := time.Now()
rl.State[model] = &UsageStats{DayStart: now, DayCount: 100} rl.State[model] = &UsageStats{DayStart: now, DayCount: 100}
for i := 0; i < 100; i++ { for i := range 100 {
t := now.Add(-time.Duration(i) * time.Second) t := now.Add(-time.Duration(i) * time.Second)
rl.State[model].Requests = append(rl.State[model].Requests, t) rl.State[model].Requests = append(rl.State[model].Requests, t)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100}) rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100})