blob: a18ce3c1bd28cea8f3002f76fc53bf7398144a6c (
plain) (
blame)
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
|
package rateio
import (
"errors"
"time"
)
const minInt = -int(^uint(0)>>1) - 1
// ErrRateExceeded is the error returned when the read rate exceeds our specification.
var ErrRateExceeded = errors.New("Read rate exceeded.")
// Limiter is an interface for a rate limiter.
// There are a few example limiters included in the package, but feel free to go wild with your own.
type Limiter interface {
// Apply this many bytes to the limiter, return ErrRateExceeded if the defined rate is exceeded.
Count(int) error
}
// simpleLimiter is a rate limiter that restricts Amount bytes in Frequency duration.
type simpleLimiter struct {
Amount int
Frequency time.Duration
numRead int
timeRead time.Time
}
// NewSimpleLimiter creates a Limiter that restricts a given number of bytes per frequency.
func NewSimpleLimiter(amount int, frequency time.Duration) *simpleLimiter {
return &simpleLimiter{
Amount: amount,
Frequency: frequency,
}
}
// NewGracefulLimiter returns a Limiter that is the same as a
// SimpleLimiter but adds a grace period at the start of the rate
// limiting where it allows unlimited bytes to be read during that
// period.
func NewGracefulLimiter(amount int, frequency time.Duration, grace time.Duration) *simpleLimiter {
return &simpleLimiter{
Amount: amount,
Frequency: frequency,
numRead: minInt,
timeRead: time.Now().Add(grace),
}
}
// Count applies n bytes to the limiter.
func (limit *simpleLimiter) Count(n int) error {
now := time.Now()
if now.After(limit.timeRead) {
limit.numRead = 0
limit.timeRead = now.Add(limit.Frequency)
}
limit.numRead += n
if limit.numRead > limit.Amount {
return ErrRateExceeded
}
return nil
}
// Delay returns a channel that can be used to block until next window
func (limit *simpleLimiter) Delay() <-chan time.Time {
waitTill := time.Now()
if limit.numRead >= limit.Amount {
waitTill = waitTill.Add(limit.Frequency)
}
return time.NewTimer(time.Until(waitTill)).C
}
|
