From 3d3dbc8fbfac21a75e1197832fb191a182feffbb Mon Sep 17 00:00:00 2001
From: Konstantin Maksimov <kmaksimov@gmail.com>
Date: Mon, 1 Jul 2019 19:19:40 +0300
Subject: [PATCH] Handle RequestTimeout and TooManyRequests

Retry 5 times in case of StatusRequestTimeout StatusTooManyRequests
This fixes the issue #2680 "Make content fetch retry more robust"

Signed-off-by: Konstantin Maksimov <kmaksimov@gmail.com>
---
 remotes/docker/fetcher_test.go | 21 +++++++++++++++++++++
 remotes/docker/resolver.go     | 10 ++++++----
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/remotes/docker/fetcher_test.go b/remotes/docker/fetcher_test.go
index b6a6639ea..22fe8c181 100644
--- a/remotes/docker/fetcher_test.go
+++ b/remotes/docker/fetcher_test.go
@@ -109,6 +109,7 @@ func TestDockerFetcherOpen(t *testing.T) {
 		wantErr                bool
 		wantServerMessageError bool
 		wantPlainError         bool
+		retries                int
 	}{
 		{
 			name:         "should return status and error.message if it exists if the registry request fails",
@@ -128,12 +129,31 @@ func TestDockerFetcherOpen(t *testing.T) {
 			want:           nil,
 			wantErr:        true,
 			wantPlainError: true,
+		}, {
+			name:           "should return StatusRequestTimeout after 5 retries",
+			mockedStatus:   http.StatusRequestTimeout,
+			mockedErr:      fmt.Errorf(http.StatusText(http.StatusRequestTimeout)),
+			want:           nil,
+			wantErr:        true,
+			wantPlainError: true,
+			retries:        5,
+		}, {
+			name:           "should return StatusTooManyRequests after 5 retries",
+			mockedStatus:   http.StatusTooManyRequests,
+			mockedErr:      fmt.Errorf(http.StatusText(http.StatusTooManyRequests)),
+			want:           nil,
+			wantErr:        true,
+			wantPlainError: true,
+			retries:        5,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 
 			s := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+				if tt.retries > 0 {
+					tt.retries--
+				}
 				rw.WriteHeader(tt.mockedStatus)
 				bytes, _ := json.Marshal(tt.mockedErr)
 				rw.Write(bytes)
@@ -147,6 +167,7 @@ func TestDockerFetcherOpen(t *testing.T) {
 			got, err := r.open(context.TODO(), s.URL, "", 0)
 			assert.Equal(t, tt.wantErr, (err != nil))
 			assert.Equal(t, tt.want, got)
+			assert.Equal(t, tt.retries, 0)
 			if tt.wantErr {
 				var expectedError error
 				if tt.wantServerMessageError {
diff --git a/remotes/docker/resolver.go b/remotes/docker/resolver.go
index 45282fb51..74d1f4212 100644
--- a/remotes/docker/resolver.go
+++ b/remotes/docker/resolver.go
@@ -455,7 +455,8 @@ func (r *dockerBase) retryRequest(ctx context.Context, req *http.Request, respon
 		return nil, nil
 	}
 	last := responses[len(responses)-1]
-	if last.StatusCode == http.StatusUnauthorized {
+	switch last.StatusCode {
+	case http.StatusUnauthorized:
 		log.G(ctx).WithField("header", last.Header.Get("WWW-Authenticate")).Debug("Unauthorized")
 		if r.auth != nil {
 			if err := r.auth.AddResponses(ctx, responses); err == nil {
@@ -464,16 +465,17 @@ func (r *dockerBase) retryRequest(ctx context.Context, req *http.Request, respon
 				return nil, err
 			}
 		}
-
 		return nil, nil
-	} else if last.StatusCode == http.StatusMethodNotAllowed && req.Method == http.MethodHead {
+	case http.StatusMethodNotAllowed:
 		// Support registries which have not properly implemented the HEAD method for
 		// manifests endpoint
-		if strings.Contains(req.URL.Path, "/manifests/") {
+		if req.Method == http.MethodHead && strings.Contains(req.URL.Path, "/manifests/") {
 			// TODO: copy request?
 			req.Method = http.MethodGet
 			return copyRequest(req)
 		}
+	case http.StatusRequestTimeout, http.StatusTooManyRequests:
+		return copyRequest(req)
 	}
 
 	// TODO: Handle 50x errors accounting for attempt history