-
Notifications
You must be signed in to change notification settings - Fork 1.7k
/
ratelimiter.go
49 lines (42 loc) · 1.31 KB
/
ratelimiter.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
package common
import (
"errors"
"sync"
"golang.org/x/time/rate"
)
// Wrapper around Go's rate.Limiter that supports both global and a per-sender rate limiting.
type RateLimiter struct {
global *rate.Limiter
perSender map[string]*rate.Limiter
config RateLimiterConfig
mu sync.Mutex
}
type RateLimiterConfig struct {
GlobalRPS float64 `json:"globalRPS"`
GlobalBurst int `json:"globalBurst"`
PerSenderRPS float64 `json:"perSenderRPS"`
PerSenderBurst int `json:"perSenderBurst"`
}
func NewRateLimiter(config RateLimiterConfig) (*RateLimiter, error) {
if config.GlobalRPS <= 0.0 || config.PerSenderRPS <= 0.0 {
return nil, errors.New("RPS values must be positive")
}
if config.GlobalBurst <= 0 || config.PerSenderBurst <= 0 {
return nil, errors.New("burst values must be positive")
}
return &RateLimiter{
global: rate.NewLimiter(rate.Limit(config.GlobalRPS), config.GlobalBurst),
perSender: make(map[string]*rate.Limiter),
config: config,
}, nil
}
func (rl *RateLimiter) Allow(sender string) bool {
rl.mu.Lock()
senderLimiter, ok := rl.perSender[sender]
if !ok {
senderLimiter = rate.NewLimiter(rate.Limit(rl.config.PerSenderRPS), rl.config.PerSenderBurst)
rl.perSender[sender] = senderLimiter
}
rl.mu.Unlock()
return senderLimiter.Allow() && rl.global.Allow()
}