Add ability to specify response HTTP status code for Throttle middleware

vasayxtx · vasayxtx · commit 6b3da1b47199 · 2023-12-28T11:00:17.000+03:00
diff --git a/middleware/throttle.go b/middleware/throttle.go
@@ -22,6 +22,7 @@ type ThrottleOpts struct {
 	Limit          int
 	BacklogLimit   int
 	BacklogTimeout time.Duration
+	StatusCode     int
 }
 
 // Throttle is a middleware that limits number of currently processed requests
@@ -49,10 +50,16 @@ func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
 		panic("chi/middleware: Throttle expects backlogLimit to be positive")
 	}
 
+	statusCode := opts.StatusCode
+	if statusCode == 0 {
+		statusCode = http.StatusTooManyRequests
+	}
+
 	t := throttler{
 		tokens:         make(chan token, opts.Limit),
 		backlogTokens:  make(chan token, opts.Limit+opts.BacklogLimit),
 		backlogTimeout: opts.BacklogTimeout,
+		statusCode:     statusCode,
 		retryAfterFn:   opts.RetryAfterFn,
 	}
 
@@ -72,7 +79,7 @@ func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
 
 			case <-ctx.Done():
 				t.setRetryAfterHeaderIfNeeded(w, true)
-				http.Error(w, errContextCanceled, http.StatusTooManyRequests)
+				http.Error(w, errContextCanceled, t.statusCode)
 				return
 
 			case btok := <-t.backlogTokens:
@@ -85,12 +92,12 @@ func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
 				select {
 				case <-timer.C:
 					t.setRetryAfterHeaderIfNeeded(w, false)
-					http.Error(w, errTimedOut, http.StatusTooManyRequests)
+					http.Error(w, errTimedOut, t.statusCode)
 					return
 				case <-ctx.Done():
 					timer.Stop()
 					t.setRetryAfterHeaderIfNeeded(w, true)
-					http.Error(w, errContextCanceled, http.StatusTooManyRequests)
+					http.Error(w, errContextCanceled, t.statusCode)
 					return
 				case tok := <-t.tokens:
 					defer func() {
@@ -103,7 +110,7 @@ func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
 
 			default:
 				t.setRetryAfterHeaderIfNeeded(w, false)
-				http.Error(w, errCapacityExceeded, http.StatusTooManyRequests)
+				http.Error(w, errCapacityExceeded, t.statusCode)
 				return
 			}
 		}
@@ -119,8 +126,9 @@ type token struct{}
 type throttler struct {
 	tokens         chan token
 	backlogTokens  chan token
-	retryAfterFn   func(ctxDone bool) time.Duration
 	backlogTimeout time.Duration
+	statusCode     int
+	retryAfterFn   func(ctxDone bool) time.Duration
 }
 
 // setRetryAfterHeaderIfNeeded sets Retry-After HTTP header if corresponding retryAfterFn option of throttler is initialized.
diff --git a/middleware/throttle_test.go b/middleware/throttle_test.go
@@ -116,7 +116,6 @@ func TestThrottleTriggerGatewayTimeout(t *testing.T) {
 			res, err := client.Get(server.URL)
 			assertNoError(t, err)
 			assertEqual(t, http.StatusOK, res.StatusCode)
-
 		}(i)
 	}
 
@@ -136,7 +135,6 @@ func TestThrottleTriggerGatewayTimeout(t *testing.T) {
 			assertNoError(t, err)
 			assertEqual(t, http.StatusTooManyRequests, res.StatusCode)
 			assertEqual(t, errTimedOut, strings.TrimSpace(string(buf)))
-
 		}(i)
 	}
 
@@ -175,7 +173,6 @@ func TestThrottleMaximum(t *testing.T) {
 			buf, err := ioutil.ReadAll(res.Body)
 			assertNoError(t, err)
 			assertEqual(t, testContent, buf)
-
 		}(i)
 	}
 
@@ -196,7 +193,6 @@ func TestThrottleMaximum(t *testing.T) {
 			assertNoError(t, err)
 			assertEqual(t, http.StatusTooManyRequests, res.StatusCode)
 			assertEqual(t, errCapacityExceeded, strings.TrimSpace(string(buf)))
-
 		}(i)
 	}
 
@@ -252,3 +248,39 @@ func TestThrottleMaximum(t *testing.T) {
 
 	wg.Wait()
 }*/
+
+func TestThrottleCustomStatusCode(t *testing.T) {
+	block := make(chan struct{})
+
+	r := chi.NewRouter()
+	r.Use(ThrottleWithOpts(ThrottleOpts{Limit: 1, StatusCode: http.StatusServiceUnavailable}))
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		block <- struct{}{}
+		block <- struct{}{}
+		w.Write(testContent)
+	})
+	server := httptest.NewServer(r)
+	defer server.Close()
+
+	client := http.Client{
+		Timeout: time.Second * 60, // Maximum waiting time.
+	}
+
+	done := make(chan struct{})
+
+	go func() {
+		res, err := client.Get(server.URL)
+		assertNoError(t, err)
+		assertEqual(t, http.StatusOK, res.StatusCode)
+		done <- struct{}{}
+	}()
+
+	<-block
+	res, err := client.Get(server.URL)
+	assertNoError(t, err)
+	assertEqual(t, http.StatusServiceUnavailable, res.StatusCode)
+	<-block
+
+	<-done
+}
