mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
fix(utils): honor Retry-After for 429 retries
This commit is contained in:
+28
-1
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user