Package list golang-github-go-kit-kit / 2b43a61
ratelimit: add TokenBucketThrottler Peter Bourgon 6 years ago
2 changed file(s) with 96 addition(s) and 0 deletion(s). Raw diff Collapse all Expand all
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 }
0 package ratelimit_test
1
2 import (
3 "testing"
4
5 "golang.org/x/net/context"
6
7 "github.com/go-kit/kit/endpoint"
8 "github.com/go-kit/kit/ratelimit"
9 )
10
11 func TestTokenBucketThrottler(t *testing.T) {
12 e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
13 testRateLimit(t, ratelimit.NewTokenBucketThrottler(ratelimit.TokenBucketMaxRate(0))(e), 0) // all fail
14 testRateLimit(t, ratelimit.NewTokenBucketThrottler(ratelimit.TokenBucketMaxRate(1))(e), 1) // first pass
15 testRateLimit(t, ratelimit.NewTokenBucketThrottler(ratelimit.TokenBucketMaxRate(100))(e), 100) // 100 pass
16 }
17
18 func testRateLimit(t *testing.T, e endpoint.Endpoint, rate int) {
19 ctx := context.Background()
20 for i := 0; i < rate; i++ {
21 if _, err := e(ctx, struct{}{}); err != nil {
22 t.Fatalf("rate=%d: request %d/%d failed: %v", rate, i+1, rate, err)
23 }
24 }
25 if _, err := e(ctx, struct{}{}); err != ratelimit.ErrThrottled {
26 t.Errorf("rate=%d: want %v, have %v", rate, ratelimit.ErrThrottled, err)
27 }
28 }