-
Notifications
You must be signed in to change notification settings - Fork 0
/
ratelimit.go
105 lines (94 loc) · 1.84 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
94
95
96
97
98
99
100
101
102
103
104
105
package speeds
import (
"sync"
"sync/atomic"
"time"
)
type (
RateLimit struct {
MaxRate int64
starTimestamp int64
count int64
interval time.Duration
ticker *time.Ticker
muChan chan struct{}
closeChan chan struct{}
backServiceOnce sync.Once
}
// AddCountFunc func() (count int64)
)
func NewRateLimit(maxRate int64) *RateLimit {
return &RateLimit{
MaxRate: maxRate,
}
}
func (rl *RateLimit) SetInterval(i time.Duration) {
if i <= 0 {
i = 1 * time.Second
}
rl.interval = i
if rl.ticker != nil {
rl.ticker.Stop()
rl.ticker = time.NewTicker(i)
}
}
func (rl *RateLimit) Stop() {
if rl.ticker != nil {
rl.ticker.Stop()
}
if rl.closeChan != nil {
close(rl.closeChan)
}
return
}
func (rl *RateLimit) resetChan() {
if rl.muChan != nil {
close(rl.muChan)
}
rl.muChan = make(chan struct{})
}
func (rl *RateLimit) backService() {
if rl.interval <= 0 {
rl.interval = 200 * time.Millisecond
}
rl.ticker = time.NewTicker(rl.interval)
rl.closeChan = make(chan struct{})
rl.resetChan()
rl.starTimestamp = time.Now().UnixNano()
go func() {
for {
select {
case <-rl.ticker.C:
if rl.rate() <= rl.MaxRate {
rl.resetChan()
}
case <-rl.closeChan:
return
}
}
}()
}
func (rl *RateLimit) Add(count int64) {
rl.backServiceOnce.Do(rl.backService)
for {
if rl.rate() >= rl.MaxRate { // 超出最大限额
// 阻塞
<-rl.muChan
continue
}
atomic.AddInt64(&rl.count, count)
if atomic.LoadInt64(&rl.count) < 0 {
// reach the max value
atomic.StoreInt64(&rl.count, 0)
rl.starTimestamp = time.Now().Unix()
}
break
}
}
func (rl *RateLimit) rate() int64 {
timeElapseSecond := (time.Now().UnixNano() - rl.starTimestamp) / 1e9
if timeElapseSecond <= 0 {
timeElapseSecond = 1
}
return atomic.LoadInt64(&rl.count) / (timeElapseSecond)
}