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

114
ratelimit_test.go
View file

@ -46,7 +46,7 @@ func TestCanSend(t *testing.T) {
model := "test-unlimited"
rl.Quotas[model] = ModelQuota{MaxRPM: 0, MaxTPM: 0, MaxRPD: 0}
for i := 0; i < 1000; i++ {
for range 1000 {
rl.RecordUsage(model, 100, 100)
}
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"
rl.Quotas[model] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 0}
for i := 0; i < 100; i++ {
for range 100 {
rl.RecordUsage(model, 1, 1)
}
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
opsPerGoroutine := 50
for i := 0; i < goroutines; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < opsPerGoroutine; j++ {
for range goroutines {
wg.Go(func() {
for range opsPerGoroutine {
rl.CanSend(model, 10)
rl.RecordUsage(model, 5, 5)
rl.Stats(model)
}
}()
})
}
wg.Wait()
@ -682,36 +680,30 @@ func TestConcurrentResetAndRecord(t *testing.T) {
var wg sync.WaitGroup
// Writers
for i := 0; i < 5; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 100; j++ {
for range 5 {
wg.Go(func() {
for range 100 {
rl.RecordUsage(model, 1, 1)
}
}()
})
}
// Resetters
for i := 0; i < 3; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 20; j++ {
for range 3 {
wg.Go(func() {
for range 20 {
rl.Reset(model)
}
}()
})
}
// Readers
for i := 0; i < 5; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 100; j++ {
for range 5 {
wg.Go(func() {
for range 100 {
rl.AllStats()
}
}()
})
}
wg.Wait()
@ -763,7 +755,7 @@ func BenchmarkCanSend(b *testing.B) {
now := time.Now()
entries := make([]time.Time, 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
entries[i] = t
tokens[i] = TokenEntry{Time: t, Count: 10}
@ -799,7 +791,7 @@ func BenchmarkCanSendConcurrent(b *testing.B) {
// Populate window
now := time.Now()
for i := 0; i < 100; i++ {
for range 100 {
rl.State[model] = &UsageStats{DayStart: now}
rl.RecordUsage(model, 10, 10)
}
@ -1014,7 +1006,7 @@ func TestSetQuota(t *testing.T) {
rl := newTestLimiter(t)
var wg sync.WaitGroup
for i := 0; i < 10; i++ {
for i := range 10 {
wg.Add(1)
go func(n int) {
defer wg.Done()
@ -1080,7 +1072,7 @@ func TestAddProvider(t *testing.T) {
wg.Add(1)
go func(prov Provider) {
defer wg.Done()
for i := 0; i < 10; i++ {
for range 10 {
rl.AddProvider(prov)
}
}(p)
@ -1114,7 +1106,7 @@ func TestConcurrentMultipleModels(t *testing.T) {
wg.Add(1)
go func(model string) {
defer wg.Done()
for i := 0; i < iterations; i++ {
for range iterations {
rl.CanSend(model, 10)
rl.RecordUsage(model, 10, 10)
rl.Stats(model)
@ -1142,26 +1134,22 @@ func TestConcurrentPersistAndLoad(t *testing.T) {
var wg sync.WaitGroup
// Writers + persist
for g := 0; g < 3; g++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 50; i++ {
for range 3 {
wg.Go(func() {
for range 50 {
rl.RecordUsage(model, 10, 10)
_ = rl.Persist()
}
}()
})
}
// Loaders
for g := 0; g < 3; g++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 50; i++ {
for range 3 {
wg.Go(func() {
for range 50 {
_ = rl.Load()
}
}()
})
}
wg.Wait()
@ -1181,21 +1169,19 @@ func TestConcurrentAllStatsAndRecordUsage(t *testing.T) {
wg.Add(1)
go func(model string) {
defer wg.Done()
for i := 0; i < 100; i++ {
for range 100 {
rl.RecordUsage(model, 10, 10)
}
}(m)
}
// Read AllStats concurrently
for g := 0; g < 3; g++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 50; i++ {
for range 3 {
wg.Go(func() {
for range 50 {
_ = rl.AllStats()
}
}()
})
}
wg.Wait()
@ -1208,25 +1194,21 @@ func TestConcurrentWaitForCapacityAndRecordUsage(t *testing.T) {
var wg sync.WaitGroup
for g := 0; g < 5; g++ {
wg.Add(1)
go func() {
defer wg.Done()
for range 5 {
wg.Go(func() {
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
_ = rl.WaitForCapacity(ctx, model, 10)
}()
})
}
// Record usage concurrently
for g := 0; g < 5; g++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 20; i++ {
for range 5 {
wg.Go(func() {
for range 20 {
rl.RecordUsage(model, 10, 10)
}
}()
})
}
wg.Wait()
@ -1242,12 +1224,12 @@ func BenchmarkCanSendWithPrune(b *testing.B) {
// Pre-fill with a mix of old and new entries to trigger pruning
now := time.Now()
rl.State[model] = &UsageStats{DayStart: now}
for i := 0; i < 500; i++ {
for range 500 {
old := now.Add(-2 * time.Minute)
rl.State[model].Requests = append(rl.State[model].Requests, old)
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)
rl.State[model].Requests = append(rl.State[model].Requests, recent)
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()
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)
rl.State[model].Requests = append(rl.State[model].Requests, t)
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 {
rl.Quotas[m] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 100000}
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)
rl.State[m].Requests = append(rl.State[m].Requests, t)
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()
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)
rl.State[model].Requests = append(rl.State[model].Requests, t)
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}
// Record multiple usages.
for i := 0; i < 10; i++ {
for range 10 {
rl.RecordUsage(model, 50, 50)
}
@ -363,16 +363,14 @@ func TestSQLiteConcurrent_Good(t *testing.T) {
// Concurrent RecordUsage + CanSend + Persist (no Load, which would
// overwrite in-memory state and lose recordings between cycles).
for i := 0; i < goroutines; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < opsPerGoroutine; j++ {
for range goroutines {
wg.Go(func() {
for range opsPerGoroutine {
rl.RecordUsage(model, 5, 5)
rl.CanSend(model, 10)
_ = rl.Persist()
}
}()
})
}
wg.Wait()
@ -667,7 +665,7 @@ func BenchmarkSQLitePersist(b *testing.B) {
now := time.Now()
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)
rl.State[model].Requests = append(rl.State[model].Requests, t)
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()
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)
rl.State[model].Requests = append(rl.State[model].Requests, t)
rl.State[model].Tokens = append(rl.State[model].Tokens, TokenEntry{Time: t, Count: 100})