Merge pull request kubernetes#109050 from MadhavJivrajani/client-go-retry

rest: Ensure response body is fully read and closed before retry

Author: k8s-ci-robot

staging/src/k8s.io/client-go/rest/request.go

@@ -614,15 +614,14 @@ func (r *Request) Watch(ctx context.Context) (watch.Interface, error) {
 	}
 	url := r.URL().String()
 	for {
-		req, err := r.newHTTPRequest(ctx)
-		if err != nil {
-			return nil, err
-		}
-
 		if err := r.retry.Before(ctx, r); err != nil {
 			return nil, r.retry.WrapPreviousError(err)
 		}
 
+		req, err := r.newHTTPRequest(ctx)
+		if err != nil {
+			return nil, err
+		}
 		resp, err := client.Do(req)
 		updateURLMetrics(ctx, r, resp, err)
 		r.retry.After(ctx, r, resp, err)
@@ -722,18 +721,17 @@ func (r *Request) Stream(ctx context.Context) (io.ReadCloser, error) {
 
 	url := r.URL().String()
 	for {
+		if err := r.retry.Before(ctx, r); err != nil {
+			return nil, err
+		}
+
 		req, err := r.newHTTPRequest(ctx)
 		if err != nil {
 			return nil, err
 		}
 		if r.body != nil {
 			req.Body = ioutil.NopCloser(r.body)
 		}
-
-		if err := r.retry.Before(ctx, r); err != nil {
-			return nil, err
-		}
-
 		resp, err := client.Do(req)
 		updateURLMetrics(ctx, r, resp, err)
 		r.retry.After(ctx, r, resp, err)
@@ -859,14 +857,13 @@ 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.
 	for {
+		if err := r.retry.Before(ctx, r); err != nil {
+			return r.retry.WrapPreviousError(err)
+		}
 		req, err := r.newHTTPRequest(ctx)
 		if err != nil {
 			return err
 		}
-
-		if err := r.retry.Before(ctx, r); err != nil {
-			return r.retry.WrapPreviousError(err)
-		}
 		resp, err := client.Do(req)
 		updateURLMetrics(ctx, r, resp, err)
 		// The value -1 or a value of 0 with a non-nil Body indicates that the length is unknown.

staging/src/k8s.io/client-go/rest/request_test.go

@@ -938,7 +938,7 @@ func TestRequestWatch(t *testing.T) {
 			},
 			Err: true,
 			ErrFn: func(err error) bool {
-				return apierrors.IsInternalError(err)
+				return !apierrors.IsInternalError(err) && strings.Contains(err.Error(), "failed to reset the request body while retrying a request: EOF")
 			},
 		},
 		{
@@ -954,7 +954,10 @@ func TestRequestWatch(t *testing.T) {
 			serverReturns: []responseErr{
 				{response: nil, err: io.EOF},
 			},
-			Empty: true,
+			Err: true,
+			ErrFn: func(err error) bool {
+				return !apierrors.IsInternalError(err)
+			},
 		},
 		{
 			name: "max retries 2, server always returns a response with Retry-After header",
@@ -1130,7 +1133,7 @@ func TestRequestStream(t *testing.T) {
 			},
 			Err: true,
 			ErrFn: func(err error) bool {
-				return apierrors.IsInternalError(err)
+				return !apierrors.IsInternalError(err) && strings.Contains(err.Error(), "failed to reset the request body while retrying a request: EOF")
 			},
 		},
 		{
@@ -1371,8 +1374,6 @@ func (b *testBackoffManager) Sleep(d time.Duration) {
 }
 
 func TestCheckRetryClosesBody(t *testing.T) {
-	// unblock CI until http://issue.k8s.io/108906 is resolved in 1.24
-	t.Skip("http://issue.k8s.io/108906")
 	count := 0
 	ch := make(chan struct{})
 	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
@@ -2435,6 +2436,7 @@ func TestRequestWithRetry(t *testing.T) {
 		body                         io.Reader
 		serverReturns                responseErr
 		errExpected                  error
+		errContains                  string
 		transformFuncInvokedExpected int
 		roundTripInvokedExpected     int
 	}{
@@ -2451,7 +2453,7 @@ func TestRequestWithRetry(t *testing.T) {
 			body:                         &readSeeker{err: io.EOF},
 			serverReturns:                responseErr{response: retryAfterResponse(), err: nil},
 			errExpected:                  nil,
-			transformFuncInvokedExpected: 1,
+			transformFuncInvokedExpected: 0,
 			roundTripInvokedExpected:     1,
 		},
 		{
@@ -2474,7 +2476,7 @@ func TestRequestWithRetry(t *testing.T) {
 			name:                         "server returns retryable err, request body Seek returns error, retry aborted",
 			body:                         &readSeeker{err: io.EOF},
 			serverReturns:                responseErr{response: nil, err: io.ErrUnexpectedEOF},
-			errExpected:                  io.ErrUnexpectedEOF,
+			errContains:                  "failed to reset the request body while retrying a request: EOF",
 			transformFuncInvokedExpected: 0,
 			roundTripInvokedExpected:     1,
 		},
@@ -2517,8 +2519,15 @@ func TestRequestWithRetry(t *testing.T) {
 			if test.transformFuncInvokedExpected != transformFuncInvoked {
 				t.Errorf("Expected transform func to be invoked %d times, but got: %d", test.transformFuncInvokedExpected, transformFuncInvoked)
 			}
-			if test.errExpected != unWrap(err) {
-				t.Errorf("Expected error: %v, but got: %v", test.errExpected, unWrap(err))
+			switch {
+			case test.errExpected != nil:
+				if test.errExpected != unWrap(err) {
+					t.Errorf("Expected error: %v, but got: %v", test.errExpected, unWrap(err))
+				}
+			case len(test.errContains) > 0:
+				if !strings.Contains(err.Error(), test.errContains) {
+					t.Errorf("Expected error message to caontain: %q, but got: %q", test.errContains, err.Error())
+				}
 			}
 		})
 	}
@@ -3531,3 +3540,103 @@ func TestTransportConcurrency(t *testing.T) {
 		})
 	}
 }
+
+// TODO: see if we can consolidate the other trackers into one.
+type requestBodyTracker struct {
+	io.ReadSeeker
+	f func(string)
+}
+
+func (t *requestBodyTracker) Read(p []byte) (int, error) {
+	t.f("Request.Body.Read")
+	return t.ReadSeeker.Read(p)
+}
+
+func (t *requestBodyTracker) Seek(offset int64, whence int) (int64, error) {
+	t.f("Request.Body.Seek")
+	return t.ReadSeeker.Seek(offset, whence)
+}
+
+type responseBodyTracker struct {
+	io.ReadCloser
+	f func(string)
+}
+
+func (t *responseBodyTracker) Read(p []byte) (int, error) {
+	t.f("Response.Body.Read")
+	return t.ReadCloser.Read(p)
+}
+
+func (t *responseBodyTracker) Close() error {
+	t.f("Response.Body.Close")
+	return t.ReadCloser.Close()
+}
+
+type recorder struct {
+	order []string
+}
+
+func (r *recorder) record(call string) {
+	r.order = append(r.order, call)
+}
+
+func TestRequestBodyResetOrder(t *testing.T) {
+	recorder := &recorder{}
+	respBodyTracker := &responseBodyTracker{
+		ReadCloser: nil, // the server will fill it
+		f:          recorder.record,
+	}
+
+	var attempts int
+	client := clientForFunc(func(req *http.Request) (*http.Response, error) {
+		defer func() {
+			attempts++
+		}()
+
+		// read the request body.
+		ioutil.ReadAll(req.Body)
+
+		// first attempt, we send a retry-after
+		if attempts == 0 {
+			resp := retryAfterResponse()
+			respBodyTracker.ReadCloser = ioutil.NopCloser(bytes.NewReader([]byte{}))
+			resp.Body = respBodyTracker
+			return resp, nil
+		}
+
+		return &http.Response{StatusCode: http.StatusOK}, nil
+	})
+
+	reqBodyTracker := &requestBodyTracker{
+		ReadSeeker: bytes.NewReader([]byte{}), // empty body ensures one Read operation at most.
+		f:          recorder.record,
+	}
+	req := &Request{
+		verb: "POST",
+		body: reqBodyTracker,
+		c: &RESTClient{
+			content: defaultContentConfig(),
+			Client:  client,
+		},
+		backoff: &noSleepBackOff{},
+		retry:   &withRetry{maxRetries: 1},
+	}
+
+	req.Do(context.Background())
+
+	expected := []string{
+		// 1st attempt: the server handler reads the request body
+		"Request.Body.Read",
+		// the server sends a retry-after, client reads the
+		// response body, and closes it
+		"Response.Body.Read",
+		"Response.Body.Close",
+		// client retry logic seeks to the beginning of the request body
+		"Request.Body.Seek",
+		// 2nd attempt: the server reads the request body
+		"Request.Body.Read",
+	}
+	if !reflect.DeepEqual(expected, recorder.order) {
+		t.Errorf("Expected invocation request and response body operations for retry do not match: %s", cmp.Diff(expected, recorder.order))
+	}
+}

staging/src/k8s.io/client-go/rest/with_retry.go

@@ -78,8 +78,12 @@ type WithRetry interface {
 	IsNextRetry(ctx context.Context, restReq *Request, httpReq *http.Request, resp *http.Response, err error, f IsRetryableErrorFunc) bool
 
 	// Before should be invoked prior to each attempt, including
-	// the first one. if an error is returned, the request
-	// should be aborted immediately.
+	// the first one. If an error is returned, the request should
+	// be aborted immediately.
+	//
+	// Before may also be additionally responsible for preparing
+	// the request for the next retry, namely in terms of resetting
+	// the request body in case it has been read.
 	Before(ctx context.Context, r *Request) error
 
 	// After should be invoked immediately after an attempt is made.
@@ -194,46 +198,18 @@ func (r *withRetry) IsNextRetry(ctx context.Context, restReq *Request, httpReq *
 	r.retryAfter.Wait = time.Duration(seconds) * time.Second
 	r.retryAfter.Reason = getRetryReason(r.attempts, seconds, resp, err)
 
-	if err := r.prepareForNextRetry(ctx, restReq); err != nil {
-		klog.V(4).Infof("Could not retry request - %v", err)
-		return false
-	}
-
 	return true
 }
 
-// prepareForNextRetry is responsible for carrying out operations that need
-// to be completed before the next retry is initiated:
-// - if the request context is already canceled there is no need to
-//   retry, the function will return ctx.Err().
-// - we need to seek to the beginning of the request body before we
-//   initiate the next retry, the function should return an error if
-//   it fails to do so.
-func (r *withRetry) prepareForNextRetry(ctx context.Context, request *Request) error {
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
-	// Ensure the response body is fully read and closed before
-	// we reconnect, so that we reuse the same TCP connection.
-	if seeker, ok := request.body.(io.Seeker); ok && request.body != nil {
-		if _, err := seeker.Seek(0, 0); err != nil {
-			return fmt.Errorf("can't Seek() back to beginning of body for %T", request)
-		}
-	}
-
-	klog.V(4).Infof("Got a Retry-After %s response for attempt %d to %v", r.retryAfter.Wait, r.retryAfter.Attempt, request.URL().String())
-	return nil
-}
-
 func (r *withRetry) Before(ctx context.Context, request *Request) error {
+	// If the request context is already canceled there
+	// is no need to retry.
 	if ctx.Err() != nil {
 		r.trackPreviousError(ctx.Err())
 		return ctx.Err()
 	}
 
 	url := request.URL()
-
 	// r.retryAfter represents the retry after parameters calculated
 	// from the (response, err) tuple from the last attempt, so 'Before'
 	// can apply these retry after parameters prior to the next attempt.
@@ -245,6 +221,18 @@ func (r *withRetry) Before(ctx context.Context, request *Request) error {
 		return nil
 	}
 
+	// At this point we've made atleast one attempt, post which the response
+	// body should have been fully read and closed in order for it to be safe
+	// to reset the request body before we reconnect, in order for us to reuse
+	// the same TCP connection.
+	if seeker, ok := request.body.(io.Seeker); ok && request.body != nil {
+		if _, err := seeker.Seek(0, io.SeekStart); err != nil {
+			err = fmt.Errorf("failed to reset the request body while retrying a request: %v", err)
+			r.trackPreviousError(err)
+			return err
+		}
+	}
+
 	// if we are here, we have made attempt(s) al least once before.
 	if request.backoff != nil {
 		// TODO(tkashem) with default set to use exponential backoff
@@ -263,6 +251,7 @@ func (r *withRetry) Before(ctx context.Context, request *Request) error {
 		return err
 	}
 
+	klog.V(4).Infof("Got a Retry-After %s response for attempt %d to %v", r.retryAfter.Wait, r.retryAfter.Attempt, request.URL().String())
 	return nil
 }