|
0 |
package ratelimit
|
|
1 |
|
|
2 |
import (
|
|
3 |
"errors"
|
|
4 |
"time"
|
|
5 |
|
|
6 |
"github.com/tsenart/tb"
|
|
7 |
"golang.org/x/net/context"
|
|
8 |
|
|
9 |
"github.com/go-kit/kit/endpoint"
|
|
10 |
)
|
|
11 |
|
|
12 |
// ErrThrottled is returned in the request path when the rate limiter is
|
|
13 |
// triggered and the request is rejected.
|
|
14 |
var ErrThrottled = errors.New("throttled")
|
|
15 |
|
|
16 |
// NewTokenBucketThrottler returns an endpoint.Middleware that acts as a rate
|
|
17 |
// limiter based on a "token-bucket" algorithm. Requests that would exceed the
|
|
18 |
// maximum request rate are rejected with an error.
|
|
19 |
func NewTokenBucketThrottler(options ...TokenBucketOption) endpoint.Middleware {
|
|
20 |
t := tokenBucketThrottler{
|
|
21 |
freq: 100 * time.Millisecond,
|
|
22 |
key: "",
|
|
23 |
rate: 100,
|
|
24 |
take: 1,
|
|
25 |
}
|
|
26 |
for _, option := range options {
|
|
27 |
option(&t)
|
|
28 |
}
|
|
29 |
throttler := tb.NewThrottler(t.freq)
|
|
30 |
return func(next endpoint.Endpoint) endpoint.Endpoint {
|
|
31 |
return func(ctx context.Context, request interface{}) (interface{}, error) {
|
|
32 |
if throttler.Halt(t.key, t.take, t.rate) {
|
|
33 |
return nil, ErrThrottled
|
|
34 |
}
|
|
35 |
return next(ctx, request)
|
|
36 |
}
|
|
37 |
}
|
|
38 |
}
|
|
39 |
|
|
40 |
type tokenBucketThrottler struct {
|
|
41 |
freq time.Duration
|
|
42 |
key string
|
|
43 |
rate int64
|
|
44 |
take int64
|
|
45 |
}
|
|
46 |
|
|
47 |
// TokenBucketOption sets an option on the token bucket throttler.
|
|
48 |
type TokenBucketOption func(*tokenBucketThrottler)
|
|
49 |
|
|
50 |
// TokenBucketFillFrequency sets the interval at which tokens are replenished
|
|
51 |
// into the bucket. By default, it's 100 milliseconds.
|
|
52 |
func TokenBucketFillFrequency(freq time.Duration) TokenBucketOption {
|
|
53 |
return func(t *tokenBucketThrottler) { t.freq = freq }
|
|
54 |
}
|
|
55 |
|
|
56 |
// TokenBucketMaxRate sets the maximum allowed request rate.
|
|
57 |
// By default, it's 100.
|
|
58 |
func TokenBucketMaxRate(rate int64) TokenBucketOption {
|
|
59 |
return func(t *tokenBucketThrottler) { t.rate = rate }
|
|
60 |
}
|
|
61 |
|
|
62 |
// TokenBucketTake sets the number of tokens taken with each request.
|
|
63 |
// By default, it's 1.
|
|
64 |
func TokenBucketTake(take int64) TokenBucketOption {
|
|
65 |
return func(t *tokenBucketThrottler) { t.take = take }
|
|
66 |
}
|