ratelimit.gno

  1// Package ratelimit is a deterministic token-bucket rate limiter — a port of
  2// golang.org/x/time/rate with the wall clock replaced by a caller-supplied
  3// monotonic tick (on-chain, that tick is the block height).
  4//
  5// It is a pure library: it imports no chain APIs and reads no ambient state.
  6// A [Limiter] owns a set of per-key token buckets sharing one rate/burst
  7// config; the caller decides what a key is (typically an address string) and
  8// supplies the current tick on every call. Between two observations at ticks
  9// `last` and `now`, a bucket gains (now-last)*rate tokens, capped at `burst`.
 10// Everything is integer/float arithmetic over persistent avl state, hence
 11// deterministic and replayable.
 12//
 13// A realm wires it up by holding a *Limiter in a package-level var and feeding
 14// it runtime.ChainHeight() as the tick. For a complete, live example see the
 15// demo realm [r/moul/x/daily/ratelimitdemo](/r/moul/x/daily/ratelimitdemo/v0).
 16package ratelimit
 17
 18import "gno.land/p/nt/avl/v0"
 19
 20// bucket is the per-key state persisted in the avl tree.
 21type bucket struct {
 22	tokens float64 // tokens available as of `last`
 23	last   int64   // tick at which `tokens` was computed
 24}
 25
 26// Limiter is a set of per-key token buckets sharing one rate/burst config.
 27// The zero value is not usable; construct one with [New].
 28type Limiter struct {
 29	rate    float64   // tokens replenished per tick
 30	burst   float64   // bucket capacity (max tokens, largest single burst)
 31	buckets *avl.Tree // key string -> *bucket, ordered by key (deterministic)
 32}
 33
 34// New returns a Limiter replenishing `rate` tokens per tick, each bucket capped
 35// at `burst`. rate is clamped to >= 0, burst to >= 1.
 36func New(rate, burst float64) *Limiter {
 37	l := &Limiter{buckets: avl.NewTree()}
 38	l.SetConfig(rate, burst)
 39	return l
 40}
 41
 42// SetConfig updates the shared rate and burst. rate is clamped to >= 0, burst
 43// to >= 1. Existing buckets keep their stored tokens; the new config applies
 44// from the next refill.
 45func (l *Limiter) SetConfig(rate, burst float64) {
 46	if rate < 0 {
 47		rate = 0
 48	}
 49	if burst < 1 {
 50		burst = 1
 51	}
 52	l.rate = rate
 53	l.burst = burst
 54}
 55
 56// Config returns the current rate (tokens/tick) and burst (capacity).
 57func (l *Limiter) Config() (rate, burst float64) { return l.rate, l.burst }
 58
 59// Allow consumes one token for `key` at tick `now` and reports whether the
 60// request is permitted. When the bucket is empty it returns false and consumes
 61// nothing. Equivalent to AllowN(key, now, 1).
 62func (l *Limiter) Allow(key string, now int64) bool {
 63	return l.AllowN(key, now, 1)
 64}
 65
 66// AllowN consumes `n` tokens for `key` at tick `now` and reports whether the
 67// request is permitted. When fewer than `n` tokens are available it returns
 68// false and consumes nothing. A key is seen for the first time with a full
 69// bucket of `burst` tokens.
 70func (l *Limiter) AllowN(key string, now int64, n float64) bool {
 71	b := l.load(key, now)
 72	ok := b.tokens >= n
 73	if ok {
 74		b.tokens -= n
 75	}
 76	l.buckets.Set(key, b)
 77	return ok
 78}
 79
 80// Tokens is a read-only view of how many whole tokens `key` has available at
 81// tick `now`, without mutating any state.
 82func (l *Limiter) Tokens(key string, now int64) int {
 83	if !l.buckets.Has(key) {
 84		return tokensToInt(l.burst)
 85	}
 86	b := l.buckets.Get(key).(*bucket)
 87	return tokensToInt(refill(b.tokens, b.last, now, l.rate, l.burst))
 88}
 89
 90// Len returns the number of keys the limiter has seen.
 91func (l *Limiter) Len() int { return l.buckets.Size() }
 92
 93// Iterate calls fn for every known key in ascending order, passing the whole
 94// tokens available at tick `now` and the last tick the key was observed.
 95// Returning true from fn stops the iteration early; Iterate reports whether it
 96// was stopped that way.
 97func (l *Limiter) Iterate(now int64, fn func(key string, tokens int, last int64) bool) bool {
 98	return l.buckets.Iterate("", "", func(key string, value any) bool {
 99		b := value.(*bucket)
100		return fn(key, tokensToInt(refill(b.tokens, b.last, now, l.rate, l.burst)), b.last)
101	})
102}
103
104// load returns the live bucket for key, refilled to `now`, creating a full
105// bucket the first time a key is seen. The returned bucket is not yet stored;
106// callers that mutate it must Set it back.
107func (l *Limiter) load(key string, now int64) *bucket {
108	if l.buckets.Has(key) {
109		b := l.buckets.Get(key).(*bucket)
110		b.tokens = refill(b.tokens, b.last, now, l.rate, l.burst)
111		b.last = now
112		return b
113	}
114	return &bucket{tokens: l.burst, last: now}
115}
116
117// --- pure helpers (unit-tested) -------------------------------------------
118
119// fmin returns the smaller of two float64 values.
120func fmin(a, b float64) float64 {
121	if a < b {
122		return a
123	}
124	return b
125}
126
127// refill returns the token count at tick `now` given `tokens` observed at tick
128// `last`, replenishing at `r` tokens/tick up to capacity `cap`. It never
129// decreases below the stored value and never exceeds `cap`.
130func refill(tokens float64, last, now int64, r, cap float64) float64 {
131	if now <= last {
132		return fmin(tokens, cap)
133	}
134	elapsed := float64(now - last)
135	return fmin(cap, tokens+elapsed*r)
136}
137
138// tokensToInt floors a token count to a whole, non-negative token.
139func tokensToInt(t float64) int {
140	if t < 0 {
141		return 0
142	}
143	return int(t)
144}