Send 'Retry-After' to comply with RFC6585.

vendor/github.com/vulcand/oxy/ratelimit/tokenlimiter.go

@@ -1,4 +1,4 @@
-// Tokenbucket based request rate limiter
+// Package ratelimit Tokenbucket based request rate limiter
 package ratelimit
 
 import (
@@ -13,6 +13,7 @@ import (
 	"github.com/vulcand/oxy/utils"
 )
 
+// DefaultCapacity default capacity
 const DefaultCapacity = 65536
 
 // RateSet maintains a set of rates. It can contain only one rate per period at a time.
@@ -31,15 +32,15 @@ func NewRateSet() *RateSet {
 // set then the new rate overrides the old one.
 func (rs *RateSet) Add(period time.Duration, average int64, burst int64) error {
 	if period <= 0 {
-		return fmt.Errorf("Invalid period: %v", period)
+		return fmt.Errorf("invalid period: %v", period)
 	}
 	if average <= 0 {
-		return fmt.Errorf("Invalid average: %v", average)
+		return fmt.Errorf("invalid average: %v", average)
 	}
 	if burst <= 0 {
-		return fmt.Errorf("Invalid burst: %v", burst)
+		return fmt.Errorf("invalid burst: %v", burst)
 	}
-	rs.m[period] = &rate{period, average, burst}
+	rs.m[period] = &rate{period: period, average: average, burst: burst}
 	return nil
 }
 
@@ -47,12 +48,15 @@ func (rs *RateSet) String() string {
 	return fmt.Sprint(rs.m)
 }
 
+// RateExtractor rate extractor
 type RateExtractor interface {
 	Extract(r *http.Request) (*RateSet, error)
 }
 
+// RateExtractorFunc rate extractor function type
 type RateExtractorFunc func(r *http.Request) (*RateSet, error)
 
+// Extract extract from request
 func (e RateExtractorFunc) Extract(r *http.Request) (*RateSet, error) {
 	return e(r)
 }
@@ -68,20 +72,24 @@ type TokenLimiter struct {
 	errHandler   utils.ErrorHandler
 	capacity     int
 	next         http.Handler
+
+	log *log.Logger
 }
 
 // New constructs a `TokenLimiter` middleware instance.
 func New(next http.Handler, extract utils.SourceExtractor, defaultRates *RateSet, opts ...TokenLimiterOption) (*TokenLimiter, error) {
 	if defaultRates == nil || len(defaultRates.m) == 0 {
-		return nil, fmt.Errorf("Provide default rates")
+		return nil, fmt.Errorf("provide default rates")
 	}
 	if extract == nil {
-		return nil, fmt.Errorf("Provide extract function")
+		return nil, fmt.Errorf("provide extract function")
 	}
 	tl := &TokenLimiter{
 		next:         next,
 		defaultRates: defaultRates,
 		extract:      extract,
+
+		log: log.StandardLogger(),
 	}
 
 	for _, o := range opts {
@@ -98,6 +106,17 @@ func New(next http.Handler, extract utils.SourceExtractor, defaultRates *RateSet
 	return tl, nil
 }
 
+// Logger defines the logger the token limiter will use.
+//
+// It defaults to logrus.StandardLogger(), the global logger used by logrus.
+func Logger(l *log.Logger) TokenLimiterOption {
+	return func(tl *TokenLimiter) error {
+		tl.log = l
+		return nil
+	}
+}
+
+// Wrap sets the next handler to be called by token limiter handler.
 func (tl *TokenLimiter) Wrap(next http.Handler) {
 	tl.next = next
 }
@@ -110,7 +129,7 @@ func (tl *TokenLimiter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	}
 
 	if err := tl.consumeRates(req, source, amount); err != nil {
-		log.Warnf("limiting request %v %v, limit: %v", req.Method, req.URL, err)
+		tl.log.Warnf("limiting request %v %v, limit: %v", req.Method, req.URL, err)
 		tl.errHandler.ServeHTTP(w, req, err)
 		return
 	}
@@ -155,7 +174,7 @@ func (tl *TokenLimiter) resolveRates(req *http.Request) *RateSet {
 
 	rates, err := tl.extractRates.Extract(req)
 	if err != nil {
-		log.Errorf("Failed to retrieve rates: %v", err)
+		tl.log.Errorf("Failed to retrieve rates: %v", err)
 		return tl.defaultRates
 	}
 
@@ -167,6 +186,7 @@ func (tl *TokenLimiter) resolveRates(req *http.Request) *RateSet {
 	return rates
 }
 
+// MaxRateError max rate error
 type MaxRateError struct {
 	delay time.Duration
 }
@@ -175,19 +195,21 @@ func (m *MaxRateError) Error() string {
 	return fmt.Sprintf("max rate reached: retry-in %v", m.delay)
 }
 
-type RateErrHandler struct {
-}
+// RateErrHandler error handler
+type RateErrHandler struct{}
 
 func (e *RateErrHandler) ServeHTTP(w http.ResponseWriter, req *http.Request, err error) {
 	if rerr, ok := err.(*MaxRateError); ok {
+		w.Header().Set("Retry-After", fmt.Sprintf("%.0f", rerr.delay.Seconds()))
 		w.Header().Set("X-Retry-In", rerr.delay.String())
-		w.WriteHeader(429)
+		w.WriteHeader(http.StatusTooManyRequests)
 		w.Write([]byte(err.Error()))
 		return
 	}
 	utils.DefaultHandler.ServeHTTP(w, req, err)
 }
 
+// TokenLimiterOption token limiter option type
 type TokenLimiterOption func(l *TokenLimiter) error
 
 // ErrorHandler sets error handler of the server
@@ -198,6 +220,7 @@ func ErrorHandler(h utils.ErrorHandler) TokenLimiterOption {
 	}
 }
 
+// ExtractRates sets the rate extractor
 func ExtractRates(e RateExtractor) TokenLimiterOption {
 	return func(cl *TokenLimiter) error {
 		cl.extractRates = e
@@ -205,6 +228,7 @@ func ExtractRates(e RateExtractor) TokenLimiterOption {
 	}
 }
 
+// Clock sets the clock
 func Clock(clock timetools.TimeProvider) TokenLimiterOption {
 	return func(cl *TokenLimiter) error {
 		cl.clock = clock
@@ -212,6 +236,7 @@ func Clock(clock timetools.TimeProvider) TokenLimiterOption {
 	}
 }
 
+// Capacity sets the capacity
 func Capacity(cap int) TokenLimiterOption {
 	return func(cl *TokenLimiter) error {
 		if cap <= 0 {