-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
limiter.go
202 lines (181 loc) · 6.54 KB
/
limiter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
// Copyright (c) Tailscale Inc & AUTHORS
// SPDX-License-Identifier: BSD-3-Clause
package limiter
import (
"fmt"
"html"
"io"
"sync"
"time"
"tailscale.com/util/lru"
)
// Limiter is a keyed token bucket rate limiter.
//
// Each key gets its own separate token bucket to pull from, enabling
// enforcement on things like "requests per IP address". To avoid
// unbounded memory growth, Limiter actually only tracks limits
// precisely for the N most recently seen keys, and assumes that
// untracked keys are well-behaved. This trades off absolute precision
// for bounded memory use, while still enforcing well for outlier
// keys.
//
// As such, Limiter should only be used in situations where "rough"
// enforcement of outliers only is sufficient, such as throttling
// egregious outlier keys (e.g. something sending 100 queries per
// second, where everyone else is sending at most 5).
//
// Each key's token bucket behaves like a regular token bucket, with
// the added feature that a bucket's token count can optionally go
// negative. This implements a form of "cooldown" for keys that exceed
// the rate limit: once a key starts getting denied, it must stop
// requesting tokens long enough for the bucket to return to a
// positive balance. If the key keeps hammering the limiter in excess
// of the rate limit, the token count will remain negative, and the
// key will not be allowed to proceed at all. This is in contrast to
// the classic token bucket, where a key trying to use more than the
// rate limit will get capped at the limit, but can still occasionally
// consume a token as one becomes available.
//
// The zero value is a valid limiter that rejects all requests. A
// useful limiter must specify a Size, Max and RefillInterval.
type Limiter[K comparable] struct {
// Size is the number of keys to track. Only the Size most
// recently seen keys have their limits enforced precisely, older
// keys are assumed to not be querying frequently enough to bother
// tracking.
Size int
// Max is the number of tokens available for a key to consume
// before time-based rate limiting kicks in. An unused limiter
// regains available tokens over time, up to Max tokens. A newly
// tracked key initially receives Max tokens.
Max int64
// RefillInterval is the interval at which a key regains tokens for
// use, up to Max tokens.
RefillInterval time.Duration
// Overdraft is the amount of additional tokens a key can be
// charged for when it exceeds its rate limit. Each additional
// request issued for the key charges one unit of overdraft, up to
// this limit. Overdraft tokens are refilled at the normal rate,
// and must be fully repaid before any tokens become available for
// requests.
//
// A non-zero Overdraft results in "cooldown" behavior: with a
// normal token bucket that bottoms out at zero tokens, an abusive
// key can still consume one token every RefillInterval. With a
// non-zero overdraft, a throttled key must stop requesting tokens
// entirely for a cooldown period, otherwise they remain
// perpetually in debt and cannot proceed at all.
Overdraft int64
mu sync.Mutex
cache *lru.Cache[K, *bucket]
}
// QPSInterval returns the interval between events corresponding to
// the given queries/second rate.
//
// This is a helper to be used when populating Limiter.RefillInterval.
func QPSInterval(qps float64) time.Duration {
return time.Duration(float64(time.Second) / qps)
}
type bucket struct {
cur int64 // current available tokens
lastUpdate time.Time // last timestamp at which cur was updated
}
// Allow charges the key one token (up to the overdraft limit), and
// reports whether the key can perform an action.
func (l *Limiter[K]) Allow(key K) bool {
return l.allow(key, time.Now())
}
func (l *Limiter[K]) allow(key K, now time.Time) bool {
l.mu.Lock()
defer l.mu.Unlock()
return l.allowBucketLocked(l.getBucketLocked(key, now), now)
}
func (l *Limiter[K]) getBucketLocked(key K, now time.Time) *bucket {
if l.cache == nil {
l.cache = &lru.Cache[K, *bucket]{MaxEntries: l.Size}
} else if b := l.cache.Get(key); b != nil {
return b
}
b := &bucket{
cur: l.Max,
lastUpdate: now.Truncate(l.RefillInterval),
}
l.cache.Set(key, b)
return b
}
func (l *Limiter[K]) allowBucketLocked(b *bucket, now time.Time) bool {
// Only update the bucket quota if needed to process request.
if b.cur <= 0 {
l.updateBucketLocked(b, now)
}
ret := b.cur > 0
if b.cur > -l.Overdraft {
b.cur--
}
return ret
}
func (l *Limiter[K]) updateBucketLocked(b *bucket, now time.Time) {
now = now.Truncate(l.RefillInterval)
if now.Before(b.lastUpdate) {
return
}
timeDelta := max(now.Sub(b.lastUpdate), 0)
tokenDelta := int64(timeDelta / l.RefillInterval)
b.cur = min(b.cur+tokenDelta, l.Max)
b.lastUpdate = now
}
// peekForTest returns the number of tokens for key, also reporting
// whether key was present.
func (l *Limiter[K]) tokensForTest(key K) (int64, bool) {
l.mu.Lock()
defer l.mu.Unlock()
if b, ok := l.cache.PeekOk(key); ok {
return b.cur, true
}
return 0, false
}
// DumpHTML writes the state of the limiter to the given writer,
// formatted as an HTML table. If onlyLimited is true, the output only
// lists keys that are currently being limited.
//
// DumpHTML blocks other callers of the limiter while it collects the
// state for dumping. It should not be called on large limiters
// involved in hot codepaths.
func (l *Limiter[K]) DumpHTML(w io.Writer, onlyLimited bool) {
l.dumpHTML(w, onlyLimited, time.Now())
}
func (l *Limiter[K]) dumpHTML(w io.Writer, onlyLimited bool, now time.Time) {
dump := l.collectDump(now)
io.WriteString(w, "<table><tr><th>Key</th><th>Tokens</th></tr>")
for _, line := range dump {
if onlyLimited && line.Tokens > 0 {
continue
}
kStr := html.EscapeString(fmt.Sprint(line.Key))
format := "<tr><td>%s</td><td>%d</td></tr>"
if !onlyLimited && line.Tokens <= 0 {
// Make limited entries stand out when showing
// limited+non-limited together
format = "<tr><td>%s</td><td><b>%d</b></td></tr>"
}
fmt.Fprintf(w, format, kStr, line.Tokens)
}
io.WriteString(w, "</table>")
}
// collectDump grabs a copy of the limiter state needed by DumpHTML.
func (l *Limiter[K]) collectDump(now time.Time) []dumpEntry[K] {
l.mu.Lock()
defer l.mu.Unlock()
ret := make([]dumpEntry[K], 0, l.cache.Len())
l.cache.ForEach(func(k K, v *bucket) {
l.updateBucketLocked(v, now) // so stats are accurate
ret = append(ret, dumpEntry[K]{k, v.cur})
})
return ret
}
// dumpEntry is the per-key information that DumpHTML needs to print
// limiter state.
type dumpEntry[K comparable] struct {
Key K
Tokens int64
}