vendor/github.com/go-kit/kit/ratelimit/token_bucket.go

   1 package ratelimit
   2
   3 import (
   4         "context"
   5         "errors"
   6
   7         "github.com/go-kit/kit/endpoint"
   8 )
   9
  10 // ErrLimited is returned in the request path when the rate limiter is
  11 // triggered and the request is rejected.
  12 var ErrLimited = errors.New("rate limit exceeded")
  13
  14 // Allower dictates whether or not a request is acceptable to run.
  15 // The Limiter from "golang.org/x/time/rate" already implements this interface,
  16 // one is able to use that in NewErroringLimiter without any modifications.
  17 type Allower interface {
  18         Allow() bool
  19 }
  20
  21 // NewErroringLimiter returns an endpoint.Middleware that acts as a rate
  22 // limiter. Requests that would exceed the
  23 // maximum request rate are simply rejected with an error.
  24 func NewErroringLimiter(limit Allower) endpoint.Middleware {
  25         return func(next endpoint.Endpoint) endpoint.Endpoint {
  26                 return func(ctx context.Context, request interface{}) (interface{}, error) {
  27                         if !limit.Allow() {
  28                                 return nil, ErrLimited
  29                         }
  30                         return next(ctx, request)
  31                 }
  32         }
  33 }
  34
  35 // Waiter dictates how long a request must be delayed.
  36 // The Limiter from "golang.org/x/time/rate" already implements this interface,
  37 // one is able to use that in NewDelayingLimiter without any modifications.
  38 type Waiter interface {
  39         Wait(ctx context.Context) error
  40 }
  41
  42 // NewDelayingLimiter returns an endpoint.Middleware that acts as a
  43 // request throttler. Requests that would
  44 // exceed the maximum request rate are delayed via the Waiter function
  45 func NewDelayingLimiter(limit Waiter) endpoint.Middleware {
  46         return func(next endpoint.Endpoint) endpoint.Endpoint {
  47                 return func(ctx context.Context, request interface{}) (interface{}, error) {
  48                         if err := limit.Wait(ctx); err != nil {
  49                                 return nil, err
  50                         }
  51                         return next(ctx, request)
  52                 }
  53         }
  54 }
  55
  56 // AllowerFunc is an adapter that lets a function operate as if
  57 // it implements Allower
  58 type AllowerFunc func() bool
  59
  60 // Allow makes the adapter implement Allower
  61 func (f AllowerFunc) Allow() bool {
  62         return f()
  63 }
  64
  65 // WaiterFunc is an adapter that lets a function operate as if
  66 // it implements Waiter
  67 type WaiterFunc func(ctx context.Context) error
  68
  69 // Wait makes the adapter implement Waiter
  70 func (f WaiterFunc) Wait(ctx context.Context) error {
  71         return f(ctx)
  72 }