llmgw/internal/accounting/limiter.go

package accounting

import (
	"sync"
	"time"
)

// Limiter enforces per-key rate limits.
//
// A sliding window rather than a fixed one. A fixed window lets a caller send
// its whole minute's allowance in the last second of one window and again in
// the first second of the next, so the observed rate is twice the configured
// limit exactly when the system is least able to take it.
//
// The window is approximated from two buckets, which is the usual compromise:
// exact enough that the doubling cannot happen, and bounded memory per key
// rather than a timestamp per request.
type Limiter struct {
	now func() time.Time

	mu   sync.Mutex
	keys map[string]*window
}

type window struct {
	// current and previous are counts in this minute and the one before.
	current, previous int64
	// othersCurrent and othersPrevious are the same counts on the other
	// control planes taking requests, as last heard: what makes a limit the
	// cluster's rather than each process's.
	othersCurrent, othersPrevious int64
	start                         time.Time
}

// NewLimiter builds an empty limiter.
func NewLimiter() *Limiter {
	return &Limiter{now: time.Now, keys: map[string]*window{}}
}

const windowLen = time.Minute

// Allow charges n against a key's limit and reports whether it fits.
//
// Charging and checking are one operation on purpose: doing them separately
// lets two concurrent requests both see room and both take it.
func (l *Limiter) Allow(key string, limit int, n int64) bool {
	if limit <= 0 {
		// No limit configured. Zero would be a limit of zero, which is why
		// callers pass a nil pointer rather than 0 for "unset".
		return true
	}
	now := l.now()

	l.mu.Lock()
	defer l.mu.Unlock()

	w := l.windowLocked(key, now)
	if w.rate(now)+float64(n) > float64(limit) {
		return false
	}
	w.current += n
	return true
}

// Charge records usage that was already allowed, for reconciliation.
//
// Token limits are checked against an estimate before the request runs, because
// the real figure only exists once it has. The difference is charged here, so a
// caller whose requests are consistently larger than estimated is still bounded
// — one window later than it would be with a perfect estimate.
func (l *Limiter) Charge(key string, n int64) {
	if n == 0 {
		return
	}
	now := l.now()
	l.mu.Lock()
	defer l.mu.Unlock()
	w := l.windowLocked(key, now)
	w.current += n
	if w.current < 0 {
		w.current = 0
	}
}

// Rate returns a key's current observed rate.
func (l *Limiter) Rate(key string) float64 {
	now := l.now()
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.windowLocked(key, now).rate(now)
}

func (l *Limiter) windowLocked(key string, now time.Time) *window {
	w, ok := l.keys[key]
	if !ok {
		w = &window{start: now.Truncate(windowLen)}
		l.keys[key] = w
	}
	w.roll(now)
	return w
}

// roll advances the window to now.
func (w *window) roll(now time.Time) {
	start := now.Truncate(windowLen)
	switch {
	case start.Equal(w.start):
		return
	case start.Sub(w.start) == windowLen:
		w.previous, w.current = w.current, 0
		w.othersPrevious, w.othersCurrent = w.othersCurrent, 0
	default:
		// More than one window has passed with no traffic.
		w.previous, w.current = 0, 0
		w.othersPrevious, w.othersCurrent = 0, 0
	}
	w.start = start
}

// rate weights the previous window by how much of it is still in view.
func (w *window) rate(now time.Time) float64 {
	cur, prev := w.current+w.othersCurrent, w.previous+w.othersPrevious
	elapsed := now.Sub(w.start)
	if elapsed >= windowLen {
		return float64(cur)
	}
	remaining := 1 - float64(elapsed)/float64(windowLen)
	return float64(prev)*remaining + float64(cur)
}

// Forget drops a key's window. Used when a key is revoked; keeping it would
// leak memory for every key that ever called.
func (l *Limiter) Forget(key string) {
	l.mu.Lock()
	delete(l.keys, key)
	l.mu.Unlock()
}