diff --git a/ratelimit.go b/ratelimit.go index c05d38e..49fdde7 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -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 } diff --git a/ratelimit_test.go b/ratelimit_test.go index d4112ab..33f1dd2 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -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}) diff --git a/sqlite_test.go b/sqlite_test.go index 4e7f7fe..72fadb4 100644 --- a/sqlite_test.go +++ b/sqlite_test.go @@ -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})