diff --git a/docs/development.md b/docs/development.md index 471d1dc..8d656a2 100644 --- a/docs/development.md +++ b/docs/development.md @@ -102,7 +102,7 @@ assert anything beyond absence of data races (the race detector does the work): ```go var wg sync.WaitGroup -for i := 0; i < 20; i++ { +for i := range 20 { wg.Add(1) go func() { defer wg.Done() diff --git a/ratelimit.go b/ratelimit.go index 49fdde7..d6c076a 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -6,10 +6,12 @@ import ( "encoding/json" "fmt" "io" + "iter" "maps" "net/http" "os" "path/filepath" + "slices" "sync" "time" @@ -276,24 +278,14 @@ func (rl *RateLimiter) prune(model string) { window := now.Add(-1 * time.Minute) // Prune requests - validReqs := 0 - for _, t := range stats.Requests { - if t.After(window) { - stats.Requests[validReqs] = t - validReqs++ - } - } - stats.Requests = stats.Requests[:validReqs] + stats.Requests = slices.DeleteFunc(stats.Requests, func(t time.Time) bool { + return !t.After(window) + }) // Prune tokens - validTokens := 0 - for _, t := range stats.Tokens { - if t.Time.After(window) { - stats.Tokens[validTokens] = t - validTokens++ - } - } - stats.Tokens = stats.Tokens[:validTokens] + stats.Tokens = slices.DeleteFunc(stats.Tokens, func(t TokenEntry) bool { + return !t.Time.After(window) + }) // Reset daily counter if day has passed if now.Sub(stats.DayStart) >= 24*time.Hour { @@ -412,6 +404,30 @@ type ModelStats struct { DayStart time.Time } +// Models returns an iterator over all model names tracked by the limiter. +func (rl *RateLimiter) Models() iter.Seq[string] { + return func(yield func(string) bool) { + stats := rl.AllStats() + for m := range stats { + if !yield(m) { + return + } + } + } +} + +// Iter returns an iterator over all model names and their current stats. +func (rl *RateLimiter) Iter() iter.Seq2[string, ModelStats] { + return func(yield func(string, ModelStats) bool) { + stats := rl.AllStats() + for k, v := range stats { + if !yield(k, v) { + return + } + } + } +} + // Stats returns current stats for a model. func (rl *RateLimiter) Stats(model string) ModelStats { rl.mu.Lock() @@ -460,23 +476,12 @@ func (rl *RateLimiter) AllStats() map[string]ModelStats { for m := range result { // Prune inline if s, ok := rl.State[m]; ok { - validReqs := 0 - for _, t := range s.Requests { - if t.After(window) { - s.Requests[validReqs] = t - validReqs++ - } - } - s.Requests = s.Requests[:validReqs] - - validTokens := 0 - for _, t := range s.Tokens { - if t.Time.After(window) { - s.Tokens[validTokens] = t - validTokens++ - } - } - s.Tokens = s.Tokens[:validTokens] + s.Requests = slices.DeleteFunc(s.Requests, func(t time.Time) bool { + return !t.After(window) + }) + s.Tokens = slices.DeleteFunc(s.Tokens, func(t TokenEntry) bool { + return !t.Time.After(window) + }) if now.Sub(s.DayStart) >= 24*time.Hour { s.DayStart = now diff --git a/ratelimit_test.go b/ratelimit_test.go index 33f1dd2..35a2f6e 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -768,7 +768,7 @@ func BenchmarkCanSend(b *testing.B) { } b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { rl.CanSend(model, 100) } } @@ -779,7 +779,7 @@ func BenchmarkRecordUsage(b *testing.B) { rl.Quotas[model] = ModelQuota{MaxRPM: 100000, MaxTPM: 1000000000, MaxRPD: 1000000} b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { rl.RecordUsage(model, 100, 100) } } @@ -1237,7 +1237,7 @@ func BenchmarkCanSendWithPrune(b *testing.B) { rl.State[model].DayCount = 1000 b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { rl.CanSend(model, 100) } } @@ -1256,7 +1256,7 @@ func BenchmarkStats(b *testing.B) { } b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { rl.Stats(model) } } @@ -1277,7 +1277,7 @@ func BenchmarkAllStats(b *testing.B) { } b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { rl.AllStats() } } @@ -1300,7 +1300,7 @@ func BenchmarkPersist(b *testing.B) { } b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { _ = rl.Persist() } } diff --git a/sqlite_test.go b/sqlite_test.go index 72fadb4..e2e13a6 100644 --- a/sqlite_test.go +++ b/sqlite_test.go @@ -672,7 +672,7 @@ func BenchmarkSQLitePersist(b *testing.B) { } b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { _ = rl.Persist() } } @@ -698,7 +698,7 @@ func BenchmarkSQLiteLoad(b *testing.B) { _ = rl.Persist() b.ResetTimer() - for i := 0; i < b.N; i++ { + for range b.N { _ = rl.Load() } }