ratelimit.mx raw

   1  // Package ratelimit provides a cooperative token-bucket rate limiter.
   2  // No mutexes — designed for single-threaded cooperative scheduling.
   3  package ratelimit
   4  
   5  import "time"
   6  
   7  type bucket struct {
   8  	tokens float64
   9  	last   int64 // unix nano
  10  }
  11  
  12  // Limiter is a per-key token bucket rate limiter.
  13  type Limiter struct {
  14  	buckets map[string]*bucket
  15  	rate    float64 // tokens per second
  16  	burst   int     // max tokens
  17  }
  18  
  19  // New creates a rate limiter. Rate is tokens/second, burst is the
  20  // maximum tokens that can accumulate.
  21  func New(rate float64, burst int) *Limiter {
  22  	return &Limiter{
  23  		buckets: map[string]*bucket{},
  24  		rate:    rate,
  25  		burst:   burst,
  26  	}
  27  }
  28  
  29  // Allow checks whether key has a token available. Consumes one token
  30  // if allowed.
  31  func (l *Limiter) Allow(key []byte) bool {
  32  	k := string(key)
  33  	now := time.Now().UnixNano()
  34  	b, ok := l.buckets[k]
  35  	if !ok {
  36  		b = &bucket{tokens: float64(l.burst), last: now}
  37  		l.buckets[k] = b
  38  	}
  39  	elapsed := float64(now-b.last) / 1e9
  40  	b.tokens += elapsed * l.rate
  41  	if b.tokens > float64(l.burst) {
  42  		b.tokens = float64(l.burst)
  43  	}
  44  	b.last = now
  45  	if b.tokens < 1.0 {
  46  		return false
  47  	}
  48  	b.tokens--
  49  	return true
  50  }
  51  
  52  // Cleanup removes entries older than maxAge to prevent unbounded growth.
  53  func (l *Limiter) Cleanup(maxAge time.Duration) {
  54  	cutoff := time.Now().Add(-maxAge).UnixNano()
  55  	for k, b := range l.buckets {
  56  		if b.last < cutoff {
  57  			delete(l.buckets, k)
  58  		}
  59  	}
  60  }
  61