/
sliding_window_limiter.go
70 lines (52 loc) · 1.56 KB
/
sliding_window_limiter.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
package slidingwindow
import (
"context"
"sync"
"time"
"github.com/air-go/rpc/library/limiter"
)
type slidingWindowEntry struct {
allow bool
}
var _ limiter.Entry = (*slidingWindowEntry)(nil)
func (e *slidingWindowEntry) Allow() bool {
return e.allow
}
func (e *slidingWindowEntry) Finish() {}
func (e *slidingWindowEntry) Error() error { return nil }
type slidingWindowLimiter struct {
limiters sync.Map
}
var _ limiter.Limiter = (*slidingWindowLimiter)(nil)
func NewLimiter() *slidingWindowLimiter {
return &slidingWindowLimiter{
limiters: sync.Map{},
}
}
func (l *slidingWindowLimiter) Check(ctx context.Context, r limiter.Resource) limiter.Entry {
return &slidingWindowEntry{allow: l.getLimiter(r).Allow()}
}
func (l *slidingWindowLimiter) SetLimit(ctx context.Context, r limiter.Resource) {
l.getLimiter(r).SetLimit(r.Limit)
}
func (l *slidingWindowLimiter) SetBurst(ctx context.Context, r limiter.Resource) {
l.getLimiter(r).SetWindow(l.burst2Window(r.Burst))
}
func (l *slidingWindowLimiter) SetWindow(ctx context.Context, r limiter.Resource) {}
func (l *slidingWindowLimiter) getLimiter(r limiter.Resource) (lim SlidingWindow) {
val, ok := l.limiters.Load(r.Name)
if !ok {
lim = NewSlidingWindow(r.Limit, l.burst2Window(r.Burst))
l.limiters.Store(r.Name, lim)
return
}
if lim, ok = val.(SlidingWindow); !ok {
lim = NewSlidingWindow(r.Limit, l.burst2Window(r.Burst))
l.limiters.Store(r.Name, lim)
return
}
return
}
func (l *slidingWindowLimiter) burst2Window(burst int) time.Duration {
return time.Duration(burst) * time.Millisecond
}