fix(utils): honor Retry-After for 429 retries

This commit is contained in:
Alix-007
2026-03-30 13:04:51 +08:00
parent 93f4c4a843
commit cd3f6600ca
2 changed files with 103 additions and 1 deletions
+28 -1
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"net/http"
"strconv"
"time"
)
@@ -36,7 +37,7 @@ func DoRequestWithRetry(client *http.Client, req *http.Request) (*http.Response,
}
if i < maxRetries-1 {
if err = sleepWithCtx(req.Context(), retryDelayUnit*time.Duration(i+1)); err != nil {
if err = sleepWithCtx(req.Context(), retryDelayForAttempt(resp, i)); err != nil {
if resp != nil {
resp.Body.Close()
}
@@ -47,6 +48,32 @@ func DoRequestWithRetry(client *http.Client, req *http.Request) (*http.Response,
return resp, err
}
func retryDelayForAttempt(resp *http.Response, attempt int) time.Duration {
fallback := retryDelayUnit * time.Duration(attempt+1)
if resp == nil || resp.StatusCode != http.StatusTooManyRequests {
return fallback
}
retryAfter := resp.Header.Get("Retry-After")
if retryAfter == "" {
return fallback
}
if seconds, err := strconv.Atoi(retryAfter); err == nil && seconds >= 0 {
return time.Duration(seconds) * time.Second
}
if when, err := http.ParseTime(retryAfter); err == nil {
delay := time.Until(when)
if delay < 0 {
return 0
}
return delay
}
return fallback
}
func sleepWithCtx(ctx context.Context, d time.Duration) error {
timer := time.NewTimer(d)
defer timer.Stop()
+75
View File
@@ -80,6 +80,81 @@ func TestDoRequestWithRetry(t *testing.T) {
}
}
func TestDoRequestWithRetry_RetryAfter429Honored(t *testing.T) {
retryDelayUnit = 10 * time.Millisecond
t.Cleanup(func() { retryDelayUnit = time.Second })
attempts := 0
var firstAttemptAt time.Time
var secondAttemptAt time.Time
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
attempts++
if attempts == 1 {
firstAttemptAt = time.Now()
w.Header().Set("Retry-After", "1")
w.WriteHeader(http.StatusTooManyRequests)
return
}
if attempts == 2 {
secondAttemptAt = time.Now()
}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
client := &http.Client{Timeout: 5 * time.Second}
req, err := http.NewRequest(http.MethodGet, server.URL, nil)
require.NoError(t, err)
resp, err := DoRequestWithRetry(client, req)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, http.StatusOK, resp.StatusCode)
resp.Body.Close()
require.Equal(t, 2, attempts)
assert.GreaterOrEqual(t, secondAttemptAt.Sub(firstAttemptAt), 900*time.Millisecond)
}
func TestDoRequestWithRetry_RetryAfter429InvalidFallsBack(t *testing.T) {
retryDelayUnit = 50 * time.Millisecond
t.Cleanup(func() { retryDelayUnit = time.Second })
attempts := 0
var firstAttemptAt time.Time
var secondAttemptAt time.Time
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
attempts++
if attempts == 1 {
firstAttemptAt = time.Now()
w.Header().Set("Retry-After", "invalid")
w.WriteHeader(http.StatusTooManyRequests)
return
}
if attempts == 2 {
secondAttemptAt = time.Now()
}
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
client := &http.Client{Timeout: 5 * time.Second}
req, err := http.NewRequest(http.MethodGet, server.URL, nil)
require.NoError(t, err)
resp, err := DoRequestWithRetry(client, req)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, http.StatusOK, resp.StatusCode)
resp.Body.Close()
require.Equal(t, 2, attempts)
assert.GreaterOrEqual(t, secondAttemptAt.Sub(firstAttemptAt), 45*time.Millisecond)
assert.Less(t, secondAttemptAt.Sub(firstAttemptAt), 500*time.Millisecond)
}
func TestDoRequestWithRetry_ContextCancel(t *testing.T) {
// Use a long retry delay so cancellation always hits during sleepWithCtx.
retryDelayUnit = 10 * time.Second