-
Notifications
You must be signed in to change notification settings - Fork 5
/
rate_limiter.go
58 lines (51 loc) · 1.14 KB
/
rate_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
package tos
import (
"sync"
"time"
)
const (
minRate = 1024
minCapacity = 10 * 1024
)
type defaultRateLimit struct {
rate int64
capacity int64
currentAmount int64
sync.Mutex
lastConsumeTime time.Time
}
func NewDefaultRateLimit(rate int64, capacity int64) RateLimiter {
if rate < minRate {
rate = minRate
}
if capacity < minCapacity {
capacity = minCapacity
}
return &defaultRateLimit{
rate: rate,
capacity: capacity,
lastConsumeTime: time.Now(),
currentAmount: capacity,
Mutex: sync.Mutex{},
}
}
func (d *defaultRateLimit) Acquire(want int64) (ok bool, timeToWait time.Duration) {
d.Lock()
defer d.Unlock()
if want > d.capacity {
want = d.capacity
}
increment := int64(time.Since(d.lastConsumeTime).Seconds() * float64(d.rate))
if increment+d.currentAmount > d.capacity {
d.currentAmount = d.capacity
} else {
d.currentAmount += increment
}
if want > d.currentAmount {
timeToWaitSec := float64(want-d.currentAmount) / float64(d.rate)
return false, time.Duration(timeToWaitSec * float64(time.Second))
}
d.lastConsumeTime = time.Now()
d.currentAmount -= want
return true, 0
}