-
Notifications
You must be signed in to change notification settings - Fork 2
/
limiter.go
95 lines (84 loc) · 1.64 KB
/
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
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
package slidingwindow
import (
"container/list"
"context"
"sync"
"time"
"github.com/haostudio/golinks/internal/limiter"
)
// New returns a new sliding window limiter.
func New(limit int, size time.Duration) limiter.Limiter {
return &slidingWindow{
limit: limit,
size: size,
windows: make(map[string]*window),
}
}
type slidingWindow struct {
sync.Mutex
limit int
size time.Duration
windows map[string]*window
}
func (w *slidingWindow) TryHit(ctx context.Context, key string) (
res limiter.Result, err error) {
w.Lock()
wd, ok := w.windows[key]
if !ok {
wd = &window{
list: new(list.List),
}
w.windows[key] = wd
}
w.Unlock()
res = wd.tryHit(w.limit, w.size)
return
}
type window struct {
sync.Mutex
list *list.List
}
func (w *window) tryHit(limit int, size time.Duration) limiter.Result {
w.Lock()
defer w.Unlock()
head := w.pruneAndPeek(size)
// calculate reset time
now := time.Now()
reset := now.Add(size)
if head != nil {
reset = head.Value.(time.Time).Add(size)
}
// within limit
l := w.list.Len()
if l < limit {
w.list.PushBack(now)
return limiter.Result{
Limit: limit,
Remaining: limit - l - 1,
Reset: reset.Unix(),
Reached: false,
}
}
// over the limit
return limiter.Result{
Limit: limit,
Remaining: 0,
Reset: reset.Unix(),
Reached: true,
}
}
func (w *window) pruneAndPeek(size time.Duration) *list.Element {
l := w.list.Len()
for l > 0 {
elem := w.list.Front()
expire := elem.Value.(time.Time)
if expire.Add(size).Before(time.Now()) {
// can abandon the elem
w.list.Remove(elem)
l--
continue
}
return elem
}
return nil
}