forked from libp2p/go-libp2p
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ratelimiter.go
75 lines (67 loc) · 2.01 KB
/
ratelimiter.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
package mocknet
import (
"sync"
"time"
)
// A RateLimiter is used by a link to determine how long to wait before sending
// data given a bandwidth cap.
type RateLimiter struct {
lock sync.Mutex
bandwidth float64 // bytes per nanosecond
allowance float64 // in bytes
maxAllowance float64 // in bytes
lastUpdate time.Time // when allowance was updated last
count int // number of times rate limiting was applied
duration time.Duration // total delay introduced due to rate limiting
}
// Creates a new RateLimiter with bandwidth (in bytes/sec)
func NewRateLimiter(bandwidth float64) *RateLimiter {
// convert bandwidth to bytes per nanosecond
b := bandwidth / float64(time.Second)
return &RateLimiter{
bandwidth: b,
allowance: 0,
maxAllowance: bandwidth,
lastUpdate: time.Now(),
}
}
// Changes bandwidth of a RateLimiter and resets its allowance
func (r *RateLimiter) UpdateBandwidth(bandwidth float64) {
r.lock.Lock()
defer r.lock.Unlock()
// Convert bandwidth from bytes/second to bytes/nanosecond
b := bandwidth / float64(time.Second)
r.bandwidth = b
// Reset allowance
r.allowance = 0
r.maxAllowance = bandwidth
r.lastUpdate = time.Now()
}
// Returns how long to wait before sending data with length 'dataSize' bytes
func (r *RateLimiter) Limit(dataSize int) time.Duration {
r.lock.Lock()
defer r.lock.Unlock()
// update time
var duration time.Duration = time.Duration(0)
if r.bandwidth == 0 {
return duration
}
current := time.Now()
elapsedTime := current.Sub(r.lastUpdate)
r.lastUpdate = current
allowance := r.allowance + float64(elapsedTime)*r.bandwidth
// allowance can't exceed bandwidth
if allowance > r.maxAllowance {
allowance = r.maxAllowance
}
allowance -= float64(dataSize)
if allowance < 0 {
// sleep until allowance is back to 0
duration = time.Duration(-allowance / r.bandwidth)
// rate limiting was applied, record stats
r.count++
r.duration += duration
}
r.allowance = allowance
return duration
}