/
registry.go
138 lines (117 loc) · 3.05 KB
/
registry.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
package ratelimit
import (
"context"
"net/http"
"sync"
"time"
log "github.com/sirupsen/logrus"
"github.com/zalando/skipper/net"
)
const (
DefaultMaxhits = 20
DefaultTimeWindow = 1 * time.Second
DefaultCleanInterval = 60 * time.Second
)
// Registry objects hold the active ratelimiters, ensure synchronized
// access to them, apply default settings and recycle the idle
// ratelimiters.
type Registry struct {
sync.Mutex
once sync.Once
closed bool
defaults Settings
global Settings
lookup map[Settings]*Ratelimit
swarm Swarmer
redisRing *net.RedisRingClient
}
// NewRegistry initializes a registry with the provided default settings.
func NewRegistry(settings ...Settings) *Registry {
return NewSwarmRegistry(nil, nil, settings...)
}
// NewSwarmRegistry initializes a registry with an optional swarm and
// the provided default settings. If swarm is nil, clusterRatelimits
// will be replaced by voidRatelimit, which is a noop limiter implementation.
func NewSwarmRegistry(swarm Swarmer, ro *net.RedisOptions, settings ...Settings) *Registry {
defaults := Settings{
Type: DisableRatelimit,
MaxHits: DefaultMaxhits,
TimeWindow: DefaultTimeWindow,
CleanInterval: DefaultCleanInterval,
}
if ro != nil && ro.MetricsPrefix == "" {
ro.MetricsPrefix = redisMetricsPrefix
}
r := &Registry{
once: sync.Once{},
defaults: defaults,
global: defaults,
lookup: make(map[Settings]*Ratelimit),
swarm: swarm,
redisRing: net.NewRedisRingClient(ro),
}
if ro != nil {
r.redisRing.StartMetricsCollection()
}
if len(settings) > 0 {
r.global = settings[0]
}
return r
}
// Close teardown Registry and dependent resources
func (r *Registry) Close() {
r.once.Do(func() {
r.closed = true
r.redisRing.Close()
for _, rl := range r.lookup {
rl.Close()
}
})
}
func (r *Registry) get(s Settings) *Ratelimit {
r.Lock()
defer r.Unlock()
rl, ok := r.lookup[s]
if !ok {
rl = newRatelimit(s, r.swarm, r.redisRing)
r.lookup[s] = rl
}
return rl
}
// Get returns a Ratelimit instance for provided Settings
func (r *Registry) Get(s Settings) *Ratelimit {
if s.Type == DisableRatelimit || s.Type == NoRatelimit {
return nil
}
return r.get(s)
}
// Check returns Settings used and the retry-after duration in case of
// request is ratelimitted. Otherwise return the Settings and 0. It is
// only used in the global ratelimit facility.
func (r *Registry) Check(req *http.Request) (Settings, int) {
if r == nil {
return Settings{}, 0
}
s := r.global
rlimit := r.Get(s)
switch s.Type {
case ClusterServiceRatelimit:
fallthrough
case ServiceRatelimit:
if rlimit.Allow(context.Background(), "") {
return s, 0
}
return s, rlimit.RetryAfter("")
case LocalRatelimit:
log.Warning("LocalRatelimit is deprecated, please use ClientRatelimit instead")
fallthrough
case ClusterClientRatelimit:
fallthrough
case ClientRatelimit:
ip := net.RemoteHost(req)
if !rlimit.Allow(context.Background(), ip.String()) {
return s, rlimit.RetryAfter(ip.String())
}
}
return Settings{}, 0
}