make client-side max-retry override'ble

This commit is contained in:
yue9944882 2020-03-27 16:15:46 +08:00
parent 84dc704679
commit f373697c64
2 changed files with 76 additions and 3 deletions

View File

@ -91,6 +91,7 @@ type Request struct {
rateLimiter flowcontrol.RateLimiter
backoff BackoffManager
timeout time.Duration
maxRetries int
// generic components accessible via method setters
verb string
@ -139,6 +140,7 @@ func NewRequest(c *RESTClient) *Request {
backoff: backoff,
timeout: timeout,
pathPrefix: pathPrefix,
maxRetries: 10,
}
switch {
@ -391,6 +393,18 @@ func (r *Request) Timeout(d time.Duration) *Request {
return r
}
// MaxRetries makes the request use the given integer as a ceiling of retrying upon receiving
// "Retry-After" headers and 429 status-code in the response. The default is 10 unless this
// function is specifically called with a different value.
// A zero maxRetries prevent it from doing retires and return an error immediately.
func (r *Request) MaxRetries(maxRetries int) *Request {
if maxRetries < 0 {
maxRetries = 0
}
r.maxRetries = maxRetries
return r
}
// Body makes the request use obj as the body. Optional.
// If obj is a string, try to read a file of that name.
// If obj is a []byte, send it directly.
@ -831,7 +845,6 @@ func (r *Request) request(ctx context.Context, fn func(*http.Request, *http.Resp
}
// Right now we make about ten retry attempts if we get a Retry-After response.
maxRetries := 10
retries := 0
for {
@ -894,7 +907,7 @@ func (r *Request) request(ctx context.Context, fn func(*http.Request, *http.Resp
}()
retries++
if seconds, wait := checkWait(resp); wait && retries < maxRetries {
if seconds, wait := checkWait(resp); wait && retries <= r.maxRetries {
if seeker, ok := r.body.(io.Seeker); ok && r.body != nil {
_, err := seeker.Seek(0, 0)
if err != nil {

View File

@ -1403,7 +1403,8 @@ func TestConnectionResetByPeerIsRetried(t *testing.T) {
return nil, &net.OpError{Err: syscall.ECONNRESET}
}),
},
backoff: backoff,
backoff: backoff,
maxRetries: 10,
}
// We expect two retries of "connection reset by peer" and the success.
_, err := req.Do(context.Background()).Raw()
@ -2218,3 +2219,62 @@ func TestThrottledLogger(t *testing.T) {
t.Fatalf("expected %v log messages, but got %v", e, a)
}
}
func TestRequestMaxRetries(t *testing.T) {
successAtNthCalls := 1
actualCalls := 0
retryOneTimeHandler := func(w http.ResponseWriter, req *http.Request) {
defer func() { actualCalls++ }()
if actualCalls >= successAtNthCalls {
w.WriteHeader(http.StatusOK)
return
}
w.Header().Set("Retry-After", "1")
w.WriteHeader(http.StatusTooManyRequests)
actualCalls++
}
testServer := httptest.NewServer(http.HandlerFunc(retryOneTimeHandler))
defer testServer.Close()
u, err := url.Parse(testServer.URL)
if err != nil {
t.Error(err)
}
testCases := []struct {
name string
maxRetries int
expectError bool
}{
{
name: "no retrying should fail",
maxRetries: 0,
expectError: true,
},
{
name: "1 max-retry should exactly work",
maxRetries: 1,
expectError: false,
},
{
name: "5 max-retry should work",
maxRetries: 5,
expectError: false,
},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
defer func() { actualCalls = 0 }()
_, err := NewRequestWithClient(u, "", defaultContentConfig(), testServer.Client()).
Verb("get").
MaxRetries(testCase.maxRetries).
AbsPath("/foo").
DoRaw(context.TODO())
hasError := err != nil
if testCase.expectError != hasError {
t.Error(" failed checking error")
}
})
}
}