1// Package ratelimit is a deterministic token-bucket rate limiter — a port of2// golang.org/x/time/rate with the wall clock replaced by a caller-supplied3// 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/burst7// config; the caller decides what a key is (typically an address string) and8// supplies the current tick on every call. Between two observations at ticks9// `last` and `now`, a bucket gains (now-last)*rate tokens, capped at `burst`.10// Everything is integer/float arithmetic over persistent avl state, hence11// deterministic and replayable.12//13// A realm wires it up by holding a *Limiter in a package-level var and feeding14// it runtime.ChainHeight() as the tick. For a complete, live example see the15// demo realm [r/moul/x/daily/ratelimitdemo](/r/moul/x/daily/ratelimitdemo/v0).16package ratelimit1718import "gno.land/p/nt/avl/v0"1920// 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 computed24}2526// 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 tick30 burst float64 // bucket capacity (max tokens, largest single burst)31 buckets *avl.Tree // key string -> *bucket, ordered by key (deterministic)32}3334// New returns a Limiter replenishing `rate` tokens per tick, each bucket capped35// 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 l40}4142// SetConfig updates the shared rate and burst. rate is clamped to >= 0, burst43// to >= 1. Existing buckets keep their stored tokens; the new config applies44// from the next refill.45func (l *Limiter) SetConfig(rate, burst float64) {46 if rate < 0 {47 rate = 048 }49 if burst < 1 {50 burst = 151 }52 l.rate = rate53 l.burst = burst54}5556// Config returns the current rate (tokens/tick) and burst (capacity).57func (l *Limiter) Config() (rate, burst float64) { return l.rate, l.burst }5859// Allow consumes one token for `key` at tick `now` and reports whether the60// request is permitted. When the bucket is empty it returns false and consumes61// nothing. Equivalent to AllowN(key, now, 1).62func (l *Limiter) Allow(key string, now int64) bool {63 return l.AllowN(key, now, 1)64}6566// AllowN consumes `n` tokens for `key` at tick `now` and reports whether the67// request is permitted. When fewer than `n` tokens are available it returns68// false and consumes nothing. A key is seen for the first time with a full69// 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 >= n73 if ok {74 b.tokens -= n75 }76 l.buckets.Set(key, b)77 return ok78}7980// Tokens is a read-only view of how many whole tokens `key` has available at81// 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}8990// Len returns the number of keys the limiter has seen.91func (l *Limiter) Len() int { return l.buckets.Size() }9293// Iterate calls fn for every known key in ascending order, passing the whole94// 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 it96// 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}103104// load returns the live bucket for key, refilled to `now`, creating a full105// 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 = now112 return b113 }114 return &bucket{tokens: l.burst, last: now}115}116117// --- pure helpers (unit-tested) -------------------------------------------118119// fmin returns the smaller of two float64 values.120func fmin(a, b float64) float64 {121 if a < b {122 return a123 }124 return b125}126127// refill returns the token count at tick `now` given `tokens` observed at tick128// `last`, replenishing at `r` tokens/tick up to capacity `cap`. It never129// 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}137138// tokensToInt floors a token count to a whole, non-negative token.139func tokensToInt(t float64) int {140 if t < 0 {141 return 0142 }143 return int(t)144}145Signatures reconstructed verbatim from vm/qfuncs — interface params keep their inline definitions.