tokenbucket.gno

  1package tokenbucket
  2
  3import (
  4	"errors"
  5
  6	u256 "gno.land/p/onbloc/math/uint256"
  7)
  8
  9var (
 10	ErrZeroCapacity         = errors.New("tokenbucket: capacity must not be zero")
 11	ErrZeroRefillRate       = errors.New("tokenbucket: refill rate must not be zero")
 12	ErrRateLimitExceeded    = errors.New("tokenbucket: rate limit exceeded")
 13	ErrTokenBucketUnderflow = errors.New("tokenbucket: token bucket underflow")
 14)
 15
 16type TokenBucket struct {
 17	Capacity   *u256.Uint
 18	Available  *u256.Uint
 19	RefillRate *u256.Uint
 20	LastRefill *u256.Uint
 21}
 22
 23func New(capacity *u256.Uint, refillRate *u256.Uint, now *u256.Uint) (*TokenBucket, error) {
 24	if capacity.IsZero() {
 25		return nil, ErrZeroCapacity
 26	}
 27
 28	if refillRate.IsZero() {
 29		return nil, ErrZeroRefillRate
 30	}
 31
 32	return &TokenBucket{
 33		Capacity:   capacity.Clone(),
 34		Available:  capacity.Clone(),
 35		RefillRate: refillRate.Clone(),
 36		LastRefill: now.Clone(),
 37	}, nil
 38}
 39
 40func (tb *TokenBucket) Refill(now *u256.Uint) {
 41	if now.Lte(tb.LastRefill) {
 42		return
 43	}
 44
 45	if tb.Available.Gte(tb.Capacity) {
 46		tb.LastRefill.Set(now)
 47		return
 48	}
 49
 50	elapsed, underflow := u256.Zero().SubOverflow(now, tb.LastRefill)
 51	if underflow {
 52		panic(ErrTokenBucketUnderflow)
 53	}
 54
 55	toRefill := saturatingMul(tb.RefillRate, elapsed)
 56	tb.Available = min(tb.Capacity, saturatingAdd(tb.Available, toRefill))
 57	tb.LastRefill.Set(now)
 58}
 59
 60func (tb *TokenBucket) RateLimit(amount *u256.Uint, now *u256.Uint) error {
 61	if amount.IsZero() {
 62		return nil
 63	}
 64
 65	tb.Refill(now)
 66
 67	if tb.Available.Lt(amount) {
 68		return ErrRateLimitExceeded
 69	}
 70
 71	tb.Available = u256.Zero().Sub(tb.Available, amount)
 72
 73	return nil
 74}
 75
 76func (tb *TokenBucket) Update(capacity *u256.Uint, refillRate *u256.Uint, reset bool) error {
 77	if capacity.IsZero() {
 78		return ErrZeroCapacity
 79	}
 80
 81	if refillRate.IsZero() {
 82		return ErrZeroRefillRate
 83	}
 84
 85	tb.Capacity = capacity.Clone()
 86	tb.RefillRate = refillRate.Clone()
 87
 88	// Mirror union's token_bucket update (token_bucket.rs L84-88): only reset
 89	// available on first init or explicit reset. Do NOT clamp available down when
 90	// capacity shrinks — union lets a now-excess available drain naturally.
 91	if tb.LastRefill.IsZero() || reset {
 92		tb.Available = tb.Capacity.Clone()
 93	}
 94
 95	return nil
 96}
 97
 98func saturatingAdd(x *u256.Uint, y *u256.Uint) *u256.Uint {
 99	z, overflow := u256.Zero().AddOverflow(x, y)
100	if overflow {
101		return u256.Zero().SetAllOne()
102	}
103
104	return z
105}
106
107func saturatingMul(x *u256.Uint, y *u256.Uint) *u256.Uint {
108	z, overflow := u256.Zero().MulOverflow(x, y)
109	if overflow {
110		return u256.Zero().SetAllOne()
111	}
112
113	return z
114}
115
116func min(x *u256.Uint, y *u256.Uint) *u256.Uint {
117	if x.Lte(y) {
118		return x.Clone()
119	}
120
121	return y.Clone()
122}