feat: modernise to Go 1.26 — slices.DeleteFunc, iterators, range
- Use slices.DeleteFunc in prune() for cleaner time-window filtering - Add Models() iter.Seq[string] and Iter() iter.Seq2[string, ModelStats] - Use range over int in benchmarks and tests - Update docs example to modern range syntax Co-Authored-By: Gemini <noreply@google.com> Co-Authored-By: Virgil <virgil@lethean.io>
This commit is contained in:
parent
852e634dc0
commit
d74811f2d0
4 changed files with 47 additions and 42 deletions
|
|
@ -102,7 +102,7 @@ assert anything beyond absence of data races (the race detector does the work):
|
||||||
|
|
||||||
```go
|
```go
|
||||||
var wg sync.WaitGroup
|
var wg sync.WaitGroup
|
||||||
for i := 0; i < 20; i++ {
|
for i := range 20 {
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func() {
|
go func() {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
|
|
|
||||||
71
ratelimit.go
71
ratelimit.go
|
|
@ -6,10 +6,12 @@ import (
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"iter"
|
||||||
"maps"
|
"maps"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"slices"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
|
@ -276,24 +278,14 @@ func (rl *RateLimiter) prune(model string) {
|
||||||
window := now.Add(-1 * time.Minute)
|
window := now.Add(-1 * time.Minute)
|
||||||
|
|
||||||
// Prune requests
|
// Prune requests
|
||||||
validReqs := 0
|
stats.Requests = slices.DeleteFunc(stats.Requests, func(t time.Time) bool {
|
||||||
for _, t := range stats.Requests {
|
return !t.After(window)
|
||||||
if t.After(window) {
|
})
|
||||||
stats.Requests[validReqs] = t
|
|
||||||
validReqs++
|
|
||||||
}
|
|
||||||
}
|
|
||||||
stats.Requests = stats.Requests[:validReqs]
|
|
||||||
|
|
||||||
// Prune tokens
|
// Prune tokens
|
||||||
validTokens := 0
|
stats.Tokens = slices.DeleteFunc(stats.Tokens, func(t TokenEntry) bool {
|
||||||
for _, t := range stats.Tokens {
|
return !t.Time.After(window)
|
||||||
if t.Time.After(window) {
|
})
|
||||||
stats.Tokens[validTokens] = t
|
|
||||||
validTokens++
|
|
||||||
}
|
|
||||||
}
|
|
||||||
stats.Tokens = stats.Tokens[:validTokens]
|
|
||||||
|
|
||||||
// Reset daily counter if day has passed
|
// Reset daily counter if day has passed
|
||||||
if now.Sub(stats.DayStart) >= 24*time.Hour {
|
if now.Sub(stats.DayStart) >= 24*time.Hour {
|
||||||
|
|
@ -412,6 +404,30 @@ type ModelStats struct {
|
||||||
DayStart time.Time
|
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.
|
// Stats returns current stats for a model.
|
||||||
func (rl *RateLimiter) Stats(model string) ModelStats {
|
func (rl *RateLimiter) Stats(model string) ModelStats {
|
||||||
rl.mu.Lock()
|
rl.mu.Lock()
|
||||||
|
|
@ -460,23 +476,12 @@ func (rl *RateLimiter) AllStats() map[string]ModelStats {
|
||||||
for m := range result {
|
for m := range result {
|
||||||
// Prune inline
|
// Prune inline
|
||||||
if s, ok := rl.State[m]; ok {
|
if s, ok := rl.State[m]; ok {
|
||||||
validReqs := 0
|
s.Requests = slices.DeleteFunc(s.Requests, func(t time.Time) bool {
|
||||||
for _, t := range s.Requests {
|
return !t.After(window)
|
||||||
if t.After(window) {
|
})
|
||||||
s.Requests[validReqs] = t
|
s.Tokens = slices.DeleteFunc(s.Tokens, func(t TokenEntry) bool {
|
||||||
validReqs++
|
return !t.Time.After(window)
|
||||||
}
|
})
|
||||||
}
|
|
||||||
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]
|
|
||||||
|
|
||||||
if now.Sub(s.DayStart) >= 24*time.Hour {
|
if now.Sub(s.DayStart) >= 24*time.Hour {
|
||||||
s.DayStart = now
|
s.DayStart = now
|
||||||
|
|
|
||||||
|
|
@ -768,7 +768,7 @@ func BenchmarkCanSend(b *testing.B) {
|
||||||
}
|
}
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
rl.CanSend(model, 100)
|
rl.CanSend(model, 100)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -779,7 +779,7 @@ func BenchmarkRecordUsage(b *testing.B) {
|
||||||
rl.Quotas[model] = ModelQuota{MaxRPM: 100000, MaxTPM: 1000000000, MaxRPD: 1000000}
|
rl.Quotas[model] = ModelQuota{MaxRPM: 100000, MaxTPM: 1000000000, MaxRPD: 1000000}
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
rl.RecordUsage(model, 100, 100)
|
rl.RecordUsage(model, 100, 100)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -1237,7 +1237,7 @@ func BenchmarkCanSendWithPrune(b *testing.B) {
|
||||||
rl.State[model].DayCount = 1000
|
rl.State[model].DayCount = 1000
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
rl.CanSend(model, 100)
|
rl.CanSend(model, 100)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -1256,7 +1256,7 @@ func BenchmarkStats(b *testing.B) {
|
||||||
}
|
}
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
rl.Stats(model)
|
rl.Stats(model)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -1277,7 +1277,7 @@ func BenchmarkAllStats(b *testing.B) {
|
||||||
}
|
}
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
rl.AllStats()
|
rl.AllStats()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -1300,7 +1300,7 @@ func BenchmarkPersist(b *testing.B) {
|
||||||
}
|
}
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
_ = rl.Persist()
|
_ = rl.Persist()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -672,7 +672,7 @@ func BenchmarkSQLitePersist(b *testing.B) {
|
||||||
}
|
}
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
_ = rl.Persist()
|
_ = rl.Persist()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -698,7 +698,7 @@ func BenchmarkSQLiteLoad(b *testing.B) {
|
||||||
_ = rl.Persist()
|
_ = rl.Persist()
|
||||||
|
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for range b.N {
|
||||||
_ = rl.Load()
|
_ = rl.Load()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue