/
rate.go
141 lines (123 loc) · 3.35 KB
/
rate.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
// Copyright 2017 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package serverutil
import (
"sync"
"time"
)
// The maximum number of visitors that a RateLimiter can track.
const rateMaxVisitors = 100000
// RateLimiter implements a rate limiter with exponential backoff,
// up to a specified maximum.
type RateLimiter struct {
// Backoff specifies an initial backoff duration for a key.
// After the first request for a given key the key will be denied until
// the backoff has passed. If another request arrives after the backoff
// but before Max, the backoff duration is doubled.
Backoff time.Duration
// Max specifies a maximum backoff duration.
Max time.Duration
mu sync.Mutex // Guards the fields below.
m map[string]*visitor
first, last *visitor
}
type visitor struct {
key string
seen time.Time
backoff time.Duration
prev, next *visitor
}
// Pass attempts to pass key through the rate limiter, returning true if key is
// within the rate limit. If it returns false it also returns the duration that
// must elapse before the key will be allowed to pass again.
func (r *RateLimiter) Pass(key string) (bool, time.Duration) {
return r.pass(time.Now(), key)
}
func (r *RateLimiter) pass(now time.Time, key string) (bool, time.Duration) {
r.mu.Lock()
defer r.mu.Unlock()
// Initialize the map lazily so that RateLimiter
// may be useful in its zero form.
if r.m == nil {
r.m = map[string]*visitor{}
}
v, ok := r.m[key]
if !ok {
// We haven't seen this visitor before,
// so add a map entry and permit it.
v = &visitor{
key: key,
seen: now,
backoff: r.Backoff,
}
r.m[key] = v
// Add visitor to the end of the list.
if r.last != nil {
r.last.next = v
v.prev = r.last
}
r.last = v
// If the list is empty, add it at the start.
if r.first == nil {
r.first = v
}
} else {
// We have seen this visitor before.
// If v.backoff has passed, permit it but double the backoff.
// Otherwise, deny it.
// If MaxBackoff has passed since its last request,
// permit it and reset the backoff to its initial state.
resetTime := v.seen.Add(r.Max)
if now.After(resetTime) {
v.backoff = r.Backoff
} else {
passTime := v.seen.Add(v.backoff)
if now.After(passTime) {
v.backoff *= 2
if v.backoff > r.Max {
v.backoff = r.Max
}
} else {
return false, passTime.Sub(now)
}
}
// Mark that we've seen this visitor now.
v.seen = now
// Move v to the end of the list, if it's not there already.
if r.last != v {
// Remove v from the list.
if v.prev != nil {
v.prev.next = v.next
} else {
r.first = v.next
}
if v.next != nil {
v.next.prev = v.prev
}
// Attach v to the end of the list.
v.prev = r.last
v.next = nil
r.last.next = v
r.last = v
}
}
// Find and delete expired visitors.
// Also check whether we have exceeded the maximum number of visitors
// that we can track at once. If so, prune back to the maximum.
drop := 0
if len(r.m) >= rateMaxVisitors {
drop = len(r.m) - rateMaxVisitors
}
for v, i := r.first, 0; v != nil; v, i = v.next, i+1 {
if !now.After(v.seen.Add(r.Max)) && i >= drop {
break
}
delete(r.m, v.key)
r.first = v.next
if v.next != nil {
v.next.prev = nil
}
}
return true, 0
}