forked from botlabs-gg/yagpdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ratelimit.go
93 lines (73 loc) · 1.38 KB
/
ratelimit.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
package reddit
import (
"sync"
"time"
)
type Ratelimiter struct {
mu sync.Mutex
Windows []*RatelimitWindow
}
func NewRatelimiter() *Ratelimiter {
return &Ratelimiter{}
}
type RatelimitWindow struct {
T time.Time
ID int64
Guilds map[int64]int
}
func (r *Ratelimiter) CheckIncrement(t time.Time, guildID int64, limit int) bool {
r.mu.Lock()
defer r.mu.Unlock()
c := 0
for _, v := range r.Windows {
if t.Sub(v.T) > time.Hour {
continue
}
c += v.Guilds[guildID]
}
if c >= limit {
return false
}
window := r.findCurrentWindow(t)
window.Guilds[guildID]++
return true
}
func (r *Ratelimiter) findCurrentWindow(t time.Time) *RatelimitWindow {
id := t.Unix() / 600
// set t to when this window ends
t = time.Unix((id*600)+600, 0)
for _, v := range r.Windows {
if v.ID == id {
return v
}
}
w := &RatelimitWindow{
T: t,
ID: id,
Guilds: make(map[int64]int),
}
r.Windows = append(r.Windows, w)
return w
}
func (r *Ratelimiter) GC(t time.Time) bool {
r.mu.Lock()
defer r.mu.Unlock()
for i, v := range r.Windows {
if t.Sub(v.T) > time.Hour {
r.Windows = append(r.Windows[:i], r.Windows[i+1:]...)
return true
}
}
return false
}
func (r *Ratelimiter) FullGC(t time.Time) {
for r.GC(t) {
}
}
func (r *Ratelimiter) RunGCLoop() {
ticker := time.NewTicker(time.Minute)
for {
<-ticker.C
r.FullGC(time.Now())
}
}