ratelimit.gno
5.41 Kb · 179 lines
1// Package ratelimit ports golang.org/x/time/rate (the standard Go
2// token-bucket rate limiter) to an on-chain gno.land realm.
3//
4// Each caller address owns a token bucket with a fixed capacity (burst)
5// that refills at a steady rate expressed in tokens per block. There is no
6// wall clock on-chain, so runtime.ChainHeight() is the monotonic clock:
7// between two observations at heights `last` and `now`, the bucket gains
8// (now-last)*rate tokens, capped at `burst`. Everything is pure, integer/
9// float arithmetic over persistent avl state, hence deterministic.
10package ratelimit
11
12import (
13 "strconv"
14
15 "chain/runtime"
16 "chain/runtime/unsafe"
17
18 "gno.land/p/nt/avl/v0"
19)
20
21// Limiter configuration, shared by every caller bucket.
22// rate = tokens replenished per block.
23// burst = bucket capacity (max tokens, also the largest single burst).
24var (
25 rate float64 = 0.5 // one token every 2 blocks
26 burst float64 = 5.0
27)
28
29// bucket is the per-caller state persisted in the avl tree.
30type bucket struct {
31 tokens float64 // tokens available as of `last`
32 last int64 // block height at which `tokens` was computed
33}
34
35// buckets maps address string -> *bucket, ordered by key (deterministic).
36var buckets = avl.NewTree()
37
38// --- pure helpers (unit-tested) -------------------------------------------
39
40// fmin returns the smaller of two float64 values.
41func fmin(a, b float64) float64 {
42 if a < b {
43 return a
44 }
45 return b
46}
47
48// refill returns the token count at height `now` given `tokens` observed at
49// height `last`, replenishing at `r` tokens/block up to capacity `cap`.
50// It never decreases below the stored value and never exceeds `cap`.
51func refill(tokens float64, last, now int64, r, cap float64) float64 {
52 if now <= last {
53 return fmin(tokens, cap)
54 }
55 elapsed := float64(now - last)
56 return fmin(cap, tokens+elapsed*r)
57}
58
59// tokensToInt floors a token count to a whole, non-negative token.
60func tokensToInt(t float64) int {
61 if t < 0 {
62 return 0
63 }
64 return int(t)
65}
66
67// --- state access ----------------------------------------------------------
68
69// load returns the live bucket for addr, refilled to `now`, creating a full
70// bucket the first time an address is seen.
71func load(key string, now int64) *bucket {
72 if buckets.Has(key) {
73 b := buckets.Get(key).(*bucket)
74 b.tokens = refill(b.tokens, b.last, now, rate, burst)
75 b.last = now
76 return b
77 }
78 return &bucket{tokens: burst, last: now}
79}
80
81// --- exported API -----------------------------------------------------------
82
83// Allow consumes one token for the calling realm/user and reports whether the
84// request is permitted. Returns false (and consumes nothing) when the bucket
85// is empty.
86func Allow(cur realm) bool {
87 addr := unsafe.PreviousRealm().Address()
88 key := string(addr)
89 now := runtime.ChainHeight()
90
91 b := load(key, now)
92 ok := b.tokens >= 1.0
93 if ok {
94 b.tokens -= 1.0
95 }
96 buckets.Set(key, b)
97 return ok
98}
99
100// SetConfig updates the shared rate (tokens per block) and burst (capacity).
101// Values are clamped to be non-negative; burst is forced to at least 1.
102func SetConfig(cur realm, newRate, newBurst float64) {
103 if newRate < 0 {
104 newRate = 0
105 }
106 if newBurst < 1 {
107 newBurst = 1
108 }
109 rate = newRate
110 burst = newBurst
111}
112
113// Tokens is a read-only view of how many whole tokens `addr` has available at
114// the current block height, without mutating any state.
115func Tokens(addr address) int {
116 key := string(addr)
117 now := runtime.ChainHeight()
118 if !buckets.Has(key) {
119 return tokensToInt(burst)
120 }
121 b := buckets.Get(key).(*bucket)
122 return tokensToInt(refill(b.tokens, b.last, now, rate, burst))
123}
124
125// --- rendering --------------------------------------------------------------
126
127func ftoa(f float64) string {
128 return strconv.FormatFloat(f, 'g', -1, 64)
129}
130
131// Render shows the limiter config and a table of known callers with their
132// current tokens and last-seen height. Path "/addr/<address>" focuses one
133// caller.
134func Render(path string) string {
135 now := runtime.ChainHeight()
136
137 if len(path) > 6 && path[:6] == "/addr/" {
138 key := path[6:]
139 out := "# Rate Limiter — caller\n\n"
140 out += "Address: `" + key + "`\n\n"
141 if !buckets.Has(key) {
142 out += "_Never seen — a full bucket of " + ftoa(burst) + " tokens is available._\n"
143 return out
144 }
145 b := buckets.Get(key).(*bucket)
146 cur := refill(b.tokens, b.last, now, rate, burst)
147 out += "- Current tokens: **" + strconv.Itoa(tokensToInt(cur)) + "** / " + ftoa(burst) + "\n"
148 out += "- Last seen at height: " + strconv.FormatInt(b.last, 10) + "\n"
149 out += "- Now (height): " + strconv.FormatInt(now, 10) + "\n"
150 return out
151 }
152
153 out := "# Rate Limiter\n\n"
154 out += "Token-bucket limiter (port of `golang.org/x/time/rate`), on-chain.\n\n"
155 out += "## Config\n\n"
156 out += "| Parameter | Value |\n|---|---|\n"
157 out += "| Rate (tokens/block) | " + ftoa(rate) + " |\n"
158 out += "| Burst (capacity) | " + ftoa(burst) + " |\n"
159 out += "| Clock height | " + strconv.FormatInt(now, 10) + " |\n\n"
160
161 out += "## Callers\n\n"
162 if buckets.Size() == 0 {
163 out += "_No caller has hit the limiter yet._\n\n"
164 } else {
165 out += "| Address | Tokens | Last height |\n|---|---|---|\n"
166 buckets.Iterate("", "", func(key string, value any) bool {
167 b := value.(*bucket)
168 cur := refill(b.tokens, b.last, now, rate, burst)
169 out += "| `" + key + "` | " + strconv.Itoa(tokensToInt(cur)) +
170 " | " + strconv.FormatInt(b.last, 10) + " |\n"
171 return false
172 })
173 out += "\n"
174 }
175
176 out += "---\n"
177 out += "Call `Allow(cur)` to consume a token; read `Tokens(addr)` for a caller's balance.\n"
178 return out
179}
