blob: ecbcc56c6ae7ae88c8cba0a3a2cf038b145e2b9b (
plain)
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
|
package ratelimit
import (
"sync"
"time"
)
type RateLimiter struct {
mu sync.Mutex
lastRequest time.Time
lastRequests []time.Time
maxRequests int
window time.Duration
minDelay time.Duration
}
func NewRateLimiter(maxRequests int, window time.Duration) *RateLimiter {
minDelay := window / time.Duration(maxRequests)
return &RateLimiter{
lastRequests: make([]time.Time, 0, maxRequests),
maxRequests: maxRequests,
window: window,
minDelay: minDelay,
}
}
func (r *RateLimiter) Wait() {
r.mu.Lock()
defer r.mu.Unlock()
now := time.Now()
if !r.lastRequest.IsZero() {
elapsed := now.Sub(r.lastRequest)
if elapsed < r.minDelay {
waitTime := r.minDelay - elapsed
r.mu.Unlock()
time.Sleep(waitTime)
r.mu.Lock()
now = time.Now()
}
}
cutoff := now.Add(-r.window)
i := 0
for i < len(r.lastRequests) && r.lastRequests[i].Before(cutoff) {
i++
}
if i > 0 {
r.lastRequests = r.lastRequests[i:]
}
if len(r.lastRequests) >= r.maxRequests {
oldestInWindow := r.lastRequests[0]
waitDuration := r.window - now.Sub(oldestInWindow)
r.mu.Unlock()
time.Sleep(waitDuration + time.Millisecond)
r.mu.Lock()
now = time.Now()
cutoff = now.Add(-r.window)
i = 0
for i < len(r.lastRequests) && r.lastRequests[i].Before(cutoff) {
i++
}
if i > 0 {
r.lastRequests = r.lastRequests[i:]
}
}
r.lastRequest = now
r.lastRequests = append(r.lastRequests, now)
}
func (r *RateLimiter) RemainingRequests() int {
r.mu.Lock()
defer r.mu.Unlock()
now := time.Now()
cutoff := now.Add(-r.window)
i := 0
for i < len(r.lastRequests) && r.lastRequests[i].Before(cutoff) {
i++
}
if i > 0 {
r.lastRequests = r.lastRequests[i:]
}
return r.maxRequests - len(r.lastRequests)
}
type MultiLimiter struct {
limiters []*RateLimiter
}
func NewMultiLimiter(limiters ...*RateLimiter) *MultiLimiter {
return &MultiLimiter{
limiters: limiters,
}
}
func (m *MultiLimiter) Wait() {
for _, limiter := range m.limiters {
limiter.Wait()
}
}
|
