llmgw/internal/auth/buckets.go

package auth

import (
	"sync"
	"time"
)

// buckets are token buckets by key: rate a second, up to burst at once. Ones
// full and unused are dropped as more arrive, so a key is not kept for ever.
type buckets struct {
	mu    sync.Mutex
	rate  float64
	burst float64
	by    map[string]*bucket
}

type bucket struct {
	tokens float64
	at     time.Time
}

func newBuckets(rate, burst float64) *buckets {
	return &buckets{rate: rate, burst: burst, by: map[string]*bucket{}}
}

// allow takes a token for key at now, reporting whether there was one.
func (b *buckets) allow(key string, now time.Time) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.by) > 50_000 {
		for k, v := range b.by {
			if now.Sub(v.at).Seconds()*b.rate >= b.burst {
				delete(b.by, k)
			}
		}
	}
	k, ok := b.by[key]
	if !ok {
		k = &bucket{tokens: b.burst, at: now}
		b.by[key] = k
	}
	k.tokens = min(b.burst, k.tokens+now.Sub(k.at).Seconds()*b.rate)
	k.at = now
	if k.tokens < 1 {
		return false
	}
	k.tokens--
	return true
}