Optimize throttle middleware by avoiding unnecessary timer creation (#1011)

vasayxtx · web-flow · commit cf537d4016bf · 2025-08-26T13:23:50.000+02:00
- Only create time.NewTimer() when actually needed to wait
- Reduces allocations from 7 to 4 and memory usage by ~54%
- Improves performance by ~50% in high-throughput scenarios
diff --git a/middleware/throttle.go b/middleware/throttle.go
@@ -83,12 +83,23 @@ func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
 				return
 
 			case btok := <-t.backlogTokens:
-				timer := time.NewTimer(t.backlogTimeout)
-
 				defer func() {
 					t.backlogTokens <- btok
 				}()
 
+				// Try to get a processing token immediately first
+				select {
+				case tok := <-t.tokens:
+					defer func() {
+						t.tokens <- tok
+					}()
+					next.ServeHTTP(w, r)
+					return
+				default:
+					// No immediate token available, need to wait with timer
+				}
+
+				timer := time.NewTimer(t.backlogTimeout)
 				select {
 				case <-timer.C:
 					t.setRetryAfterHeaderIfNeeded(w, false)
diff --git a/middleware/throttle_test.go b/middleware/throttle_test.go
@@ -311,3 +311,20 @@ func TestThrottleCustomStatusCode(t *testing.T) {
 	close(wait) // Allow the last request to proceed.
 	waitResponse(http.StatusOK)
 }
+
+func BenchmarkThrottle(b *testing.B) {
+	throttleMiddleware := ThrottleBacklog(1000, 50, time.Second)
+
+	handler := throttleMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/", nil)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		w := httptest.NewRecorder()
+		handler.ServeHTTP(w, req)
+	}
+}
