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
|
package server
import (
"sync"
"time"
)
// rateLimiter is a rate limiter for the server.
// If more than limit requests are made in the time window, the server will return 429.
type rateLimiter struct {
// limit is the rate limit to apply to the server.
limit int
// window is the window in which to apply the rate limit.
window time.Duration
// nextReset is the time at which the rate limit will reset.
nextReset time.Time
// count is the number of calls made to the server.
count int
// countLock is a mutex for the callCount.
countLock sync.Mutex
// statusCode to reply with
statusCode int
}
func newRateLimiter(limit int, window time.Duration, statusCode int) *rateLimiter {
return &rateLimiter{
limit: limit,
window: window,
statusCode: statusCode,
}
}
// exceeded checks the rate limit and returns how long to wait before the next request.
func (r *rateLimiter) exceeded() time.Duration {
r.countLock.Lock()
defer r.countLock.Unlock()
if time.Now().After(r.nextReset) {
r.count = 0
r.nextReset = time.Now().Add(r.window)
}
r.count++
if r.count > r.limit {
return time.Until(r.nextReset)
}
return 0
}
|
