chore: go mod vendor

vendor/github.com/go-chi/chi/middleware/throttle.go

@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"net/http"
+	"strconv"
 	"time"
 )
 
@@ -15,44 +16,100 @@ var (
 	defaultBacklogTimeout = time.Second * 60
 )
 
+// ThrottleOpts represents a set of throttling options.
+type ThrottleOpts struct {
+	Limit          int
+	BacklogLimit   int
+	BacklogTimeout time.Duration
+	RetryAfterFn   func(ctxDone bool) time.Duration
+}
+
 // Throttle is a middleware that limits number of currently processed requests
-// at a time.
+// at a time across all users. Note: Throttle is not a rate-limiter per user,
+// instead it just puts a ceiling on the number of currentl in-flight requests
+// being processed from the point from where the Throttle middleware is mounted.
 func Throttle(limit int) func(http.Handler) http.Handler {
-	return ThrottleBacklog(limit, 0, defaultBacklogTimeout)
+	return ThrottleWithOpts(ThrottleOpts{Limit: limit, BacklogTimeout: defaultBacklogTimeout})
 }
 
 // ThrottleBacklog is a middleware that limits number of currently processed
 // requests at a time and provides a backlog for holding a finite number of
 // pending requests.
-func ThrottleBacklog(limit int, backlogLimit int, backlogTimeout time.Duration) func(http.Handler) http.Handler {
-	if limit < 1 {
+func ThrottleBacklog(limit, backlogLimit int, backlogTimeout time.Duration) func(http.Handler) http.Handler {
+	return ThrottleWithOpts(ThrottleOpts{Limit: limit, BacklogLimit: backlogLimit, BacklogTimeout: backlogTimeout})
+}
+
+// ThrottleWithOpts is a middleware that limits number of currently processed requests using passed ThrottleOpts.
+func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
+	if opts.Limit < 1 {
 		panic("chi/middleware: Throttle expects limit > 0")
 	}
 
-	if backlogLimit < 0 {
+	if opts.BacklogLimit < 0 {
 		panic("chi/middleware: Throttle expects backlogLimit to be positive")
 	}
 
 	t := throttler{
-		tokens:         make(chan token, limit),
-		backlogTokens:  make(chan token, limit+backlogLimit),
-		backlogTimeout: backlogTimeout,
+		tokens:         make(chan token, opts.Limit),
+		backlogTokens:  make(chan token, opts.Limit+opts.BacklogLimit),
+		backlogTimeout: opts.BacklogTimeout,
+		retryAfterFn:   opts.RetryAfterFn,
 	}
 
 	// Filling tokens.
-	for i := 0; i < limit+backlogLimit; i++ {
-		if i < limit {
+	for i := 0; i < opts.Limit+opts.BacklogLimit; i++ {
+		if i < opts.Limit {
 			t.tokens <- token{}
 		}
 		t.backlogTokens <- token{}
 	}
 
-	fn := func(h http.Handler) http.Handler {
-		t.h = h
-		return &t
-	}
+	return func(next http.Handler) http.Handler {
+		fn := func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
 
-	return fn
+			select {
+
+			case <-ctx.Done():
+				t.setRetryAfterHeaderIfNeeded(w, true)
+				http.Error(w, errContextCanceled, http.StatusTooManyRequests)
+				return
+
+			case btok := <-t.backlogTokens:
+				timer := time.NewTimer(t.backlogTimeout)
+
+				defer func() {
+					t.backlogTokens <- btok
+				}()
+
+				select {
+				case <-timer.C:
+					t.setRetryAfterHeaderIfNeeded(w, false)
+					http.Error(w, errTimedOut, http.StatusTooManyRequests)
+					return
+				case <-ctx.Done():
+					timer.Stop()
+					t.setRetryAfterHeaderIfNeeded(w, true)
+					http.Error(w, errContextCanceled, http.StatusTooManyRequests)
+					return
+				case tok := <-t.tokens:
+					defer func() {
+						timer.Stop()
+						t.tokens <- tok
+					}()
+					next.ServeHTTP(w, r)
+				}
+				return
+
+			default:
+				t.setRetryAfterHeaderIfNeeded(w, false)
+				http.Error(w, errCapacityExceeded, http.StatusTooManyRequests)
+				return
+			}
+		}
+
+		return http.HandlerFunc(fn)
+	}
 }
 
 // token represents a request that is being processed.
@@ -60,42 +117,16 @@ type token struct{}
 
 // throttler limits number of currently processed requests at a time.
 type throttler struct {
-	h              http.Handler
 	tokens         chan token
 	backlogTokens  chan token
 	backlogTimeout time.Duration
+	retryAfterFn   func(ctxDone bool) time.Duration
 }
 
-// ServeHTTP is the primary throttler request handler
-func (t *throttler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
-	select {
-	case <-ctx.Done():
-		http.Error(w, errContextCanceled, http.StatusServiceUnavailable)
-		return
-	case btok := <-t.backlogTokens:
-		timer := time.NewTimer(t.backlogTimeout)
-
-		defer func() {
-			t.backlogTokens <- btok
-		}()
-
-		select {
-		case <-timer.C:
-			http.Error(w, errTimedOut, http.StatusServiceUnavailable)
-			return
-		case <-ctx.Done():
-			http.Error(w, errContextCanceled, http.StatusServiceUnavailable)
-			return
-		case tok := <-t.tokens:
-			defer func() {
-				t.tokens <- tok
-			}()
-			t.h.ServeHTTP(w, r)
-		}
-		return
-	default:
-		http.Error(w, errCapacityExceeded, http.StatusServiceUnavailable)
+// setRetryAfterHeaderIfNeeded sets Retry-After HTTP header if corresponding retryAfterFn option of throttler is initialized.
+func (t throttler) setRetryAfterHeaderIfNeeded(w http.ResponseWriter, ctxDone bool) {
+	if t.retryAfterFn == nil {
 		return
 	}
+	w.Header().Set("Retry-After", strconv.Itoa(int(t.retryAfterFn(ctxDone).Seconds())))
 }