From 3e9d5317f00e2c1e9797842b8ce9431439975af5 Mon Sep 17 00:00:00 2001
From: yue9944882 <291271447@qq.com>
Date: Fri, 27 Mar 2020 16:15:46 +0800
Subject: [PATCH] make client-side max-retry override'ble

Kubernetes-commit: f373697c64e9dc5a6618e1d360ad9a51bc8b79c7
---
 rest/request.go      | 17 ++++++++++--
 rest/request_test.go | 62 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/rest/request.go b/rest/request.go
index 1acd189e..fe048ec4 100644
--- a/rest/request.go
+++ b/rest/request.go
@@ -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 {
diff --git a/rest/request_test.go b/rest/request_test.go
index 2a183ba4..1c199160 100644
--- a/rest/request_test.go
+++ b/rest/request_test.go
@@ -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")
+			}
+		})
+	}
+}