-
Notifications
You must be signed in to change notification settings - Fork 24
/
registry.go
73 lines (57 loc) · 1.35 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
package rate
import (
"sync"
"time"
"golang.org/x/time/rate"
)
var DefaultRegistryCfx = NewRegistry()
var DefaultRegistryEth = NewRegistry()
func init() {
go DefaultRegistryCfx.gcPeriodically(5*time.Minute, 3*time.Minute)
go DefaultRegistryEth.gcPeriodically(5*time.Minute, 3*time.Minute)
}
type Registry struct {
limiters map[string]*IpLimiter
whilteList map[string]struct{}
mu sync.Mutex
}
func NewRegistry() *Registry {
return &Registry{
limiters: make(map[string]*IpLimiter),
whilteList: make(map[string]struct{}),
}
}
func (m *Registry) Get(name string) (*IpLimiter, bool) {
m.mu.Lock()
defer m.mu.Unlock()
limiter, ok := m.limiters[name]
return limiter, ok
}
func (m *Registry) WhiteListed(name string) bool {
_, existed := m.whilteList[name]
return existed
}
func (m *Registry) GetOrRegister(name string, rate rate.Limit, burst int) *IpLimiter {
m.mu.Lock()
defer m.mu.Unlock()
limiter, ok := m.limiters[name]
if !ok {
limiter = NewIpLimiter(rate, burst)
m.limiters[name] = limiter
}
return limiter
}
func (m *Registry) GC(timeout time.Duration) {
m.mu.Lock()
defer m.mu.Unlock()
for _, v := range m.limiters {
v.GC(timeout)
}
}
func (m *Registry) gcPeriodically(interval time.Duration, timeout time.Duration) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for range ticker.C {
m.GC(timeout)
}
}