token_bucket.go

Documentation: github.com/go-kit/kit/ratelimit

     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  }
    73
View as plain text