github.com/GuanceCloud/cliutils@v1.1.21/network/http/rate_limiter.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  package http
     7  
     8  import (
     9  	"net/http"
    10  	"time"
    11  
    12  	"github.com/didip/tollbooth/v6"
    13  	"github.com/didip/tollbooth/v6/limiter"
    14  )
    15  
    16  // RateLimiter used to define API request rate limiter.
    17  type RateLimiter interface {
    18  	// Detect if rate limit reached on @key
    19  	Limited(key string) bool
    20  
    21  	// Update rate limite exclusively
    22  	UpdateRate(float64)
    23  }
    24  
    25  // TODO: we should implemente a common sense rate limiter
    26  
    27  // RequestKey is a callback used to calculate request @r's ID, we can
    28  // use r.Method + tkn + r.URL.Path as the id of @r, if the ID is empty
    29  // string, it's degrade into a simple rate limiter(all API's request
    30  // are limited under the rate).
    31  type RequestKey func(r *http.Request) string
    32  
    33  // DefaultRequestKey used to get key of HTTP request if you don't know how
    34  // to get the key.
    35  func DefaultRequestKey(r *http.Request) string {
    36  	return r.Header.Get("X-Forwarded-For") + r.RemoteAddr + r.Method + r.Proto + r.URL.String()
    37  }
    38  
    39  type APIRateLimiter interface {
    40  	RequestLimited(*http.Request) bool
    41  	// If rate limited, do anything what you want(cache the request, or do nothing)
    42  	LimitReadchedCallback(*http.Request)
    43  	// Update rate limite exclusively
    44  	UpdateRate(float64)
    45  }
    46  
    47  // APIRateLimiterImpl is default implemented of APIRateLimiter based on tollbooth.
    48  type APIRateLimiterImpl struct {
    49  	*limiter.Limiter
    50  	rk RequestKey
    51  }
    52  
    53  func NewAPIRateLimiter(rate float64, rk RequestKey) *APIRateLimiterImpl {
    54  	return &APIRateLimiterImpl{
    55  		Limiter: tollbooth.NewLimiter(rate,
    56  			&limiter.ExpirableOptions{DefaultExpirationTTL: time.Second}).SetBurst(1),
    57  		rk: rk,
    58  	}
    59  }
    60  
    61  // RequestLimited used to limit query param key's value on all APIs,
    62  // it means, if @key is `token', for specific token=abc123, if limit is 100/second,
    63  // then token `abc123' can only request 100 APIs per second, no matter whiching
    64  // API the token request.
    65  func (rl *APIRateLimiterImpl) RequestLimited(r *http.Request) bool {
    66  	if rl.rk == nil {
    67  		return false // no RequestKey callback set, always pass
    68  	}
    69  
    70  	return rl.Limiter.LimitReached(rl.rk(r))
    71  }
    72  
    73  // LimitReadchedCallback do nothing, just drop the request.
    74  func (rl *APIRateLimiterImpl) LimitReadchedCallback(r *http.Request) {
    75  	// do nothing
    76  }
    77  
    78  // UpdateRate update limite rate exclusively.
    79  func (rl *APIRateLimiterImpl) UpdateRate(rate float64) {
    80  	rl.Limiter.SetMax(rate)
    81  }