forked from botlabs-gg/yagpdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
multiratelimit.go
43 lines (33 loc) · 980 Bytes
/
multiratelimit.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
package multiratelimit
import (
"golang.org/x/time/rate"
"sync"
"time"
)
type MultiRatelimiter struct {
mu sync.Mutex
limiters map[interface{}]*rate.Limiter
maxPerSecond float64
maxBurst int
}
func NewMultiRatelimiter(maxPerSecond float64, maxBurst int) *MultiRatelimiter {
multiLimiter := &MultiRatelimiter{
limiters: make(map[interface{}]*rate.Limiter),
maxPerSecond: maxPerSecond,
maxBurst: maxBurst,
}
return multiLimiter
}
func (multi *MultiRatelimiter) findCreateLimiter(key interface{}) *rate.Limiter {
multi.mu.Lock()
defer multi.mu.Unlock()
if current, ok := multi.limiters[key]; ok {
return current
}
// not found, create it
multi.limiters[key] = rate.NewLimiter(rate.Limit(multi.maxPerSecond), multi.maxBurst)
return multi.limiters[key]
}
func (multi *MultiRatelimiter) AllowN(key interface{}, now time.Time, n int) bool {
return multi.findCreateLimiter(key).AllowN(now, n)
}