func limit(b *ratelimit.Bucket, wait bool, errId string) func() error {
return func() error {
if wait {
time.Sleep(b.Take(1))
} else if b.TakeAvailable(1) == 0 {
return errors.New(errId, "too many request", 429)
}
return nil
}
}
// NewTokenBucketLimiter returns an endpoint.Middleware that acts as a rate
// limiter based on a token-bucket algorithm. Requests that would exceed the
// maximum request rate are simply rejected with an error.
func NewTokenBucketLimiter(tb *ratelimit.Bucket) endpoint.Middleware {
return func(next endpoint.Endpoint) endpoint.Endpoint {
return func() error {
if tb.TakeAvailable(1) == 0 {
return ErrLimited
}
return next()
}
}
}
func NewTokenBucketLimiter(tb *ratelimit.Bucket) endpoint.Middleware {
return func(next endpoint.Endpoint) endpoint.Endpoint {
return func(ctx context.Context, request interface{}) (interface{}, error) {
if tb.TakeAvailable(1) == 0 { // HL
return nil, ErrLimited // HL
} // HL
return next(ctx, request) // HL
}
}
}
func getTestServerThrottled(body string) *httptest.Server {
var rateLimiter *ratelimit.Bucket
// Rate limit: 2 req/s, capacity 2
rateLimiter = ratelimit.NewBucket(500*time.Millisecond, 2)
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
code := 200
if tokens := rateLimiter.TakeAvailable(1); tokens == 0 {
code = 503
}
w.WriteHeader(code)
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, body)
}))
}
