Skip to content

Commit cf537d4

Browse files
authored
Optimize throttle middleware by avoiding unnecessary timer creation (#1011)
- 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
1 parent 9040e95 commit cf537d4

File tree

2 files changed

+30
-2
lines changed

2 files changed

+30
-2
lines changed

middleware/throttle.go

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,12 +83,23 @@ func ThrottleWithOpts(opts ThrottleOpts) func(http.Handler) http.Handler {
8383
return
8484

8585
case btok := <-t.backlogTokens:
86-
timer := time.NewTimer(t.backlogTimeout)
87-
8886
defer func() {
8987
t.backlogTokens <- btok
9088
}()
9189

90+
// Try to get a processing token immediately first
91+
select {
92+
case tok := <-t.tokens:
93+
defer func() {
94+
t.tokens <- tok
95+
}()
96+
next.ServeHTTP(w, r)
97+
return
98+
default:
99+
// No immediate token available, need to wait with timer
100+
}
101+
102+
timer := time.NewTimer(t.backlogTimeout)
92103
select {
93104
case <-timer.C:
94105
t.setRetryAfterHeaderIfNeeded(w, false)

middleware/throttle_test.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -311,3 +311,20 @@ func TestThrottleCustomStatusCode(t *testing.T) {
311311
close(wait) // Allow the last request to proceed.
312312
waitResponse(http.StatusOK)
313313
}
314+
315+
func BenchmarkThrottle(b *testing.B) {
316+
throttleMiddleware := ThrottleBacklog(1000, 50, time.Second)
317+
318+
handler := throttleMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
319+
w.WriteHeader(http.StatusOK)
320+
}))
321+
322+
req := httptest.NewRequest("GET", "/", nil)
323+
324+
b.ResetTimer()
325+
b.ReportAllocs()
326+
for i := 0; i < b.N; i++ {
327+
w := httptest.NewRecorder()
328+
handler.ServeHTTP(w, req)
329+
}
330+
}

0 commit comments

Comments
 (0)