Improving retry logic in remote to handle more cases

remote/remote.go

@@ -193,14 +193,13 @@ func newV1Image(keychain authn.Keychain, repoName string, platform imgutil.Platf
 		time.Sleep(100 * time.Duration(i) * time.Millisecond) // wait if retrying
 		image, err = remote.Image(ref, remote.WithAuth(auth), remote.WithTransport(http.DefaultTransport), remote.WithPlatform(v1Platform))
 		if err != nil {
-			if err == io.EOF && i != maxRetries {
+			transportErr, ok := err.(*transport.Error)
+			isValid40x := ok && is40x(transportErr)
+			if err == io.EOF || isValid40x && i != maxRetries {
 				continue // retry if EOF
 			}
-			if transportErr, ok := err.(*transport.Error); ok && len(transportErr.Errors) > 0 {
-				switch transportErr.StatusCode {
-				case http.StatusNotFound, http.StatusUnauthorized:
-					return emptyImage(platform)
-				}
+			if isValid40x && len(transportErr.Errors) > 0 {
+				return emptyImage(platform)
 			}
 			if strings.Contains(err.Error(), "no child with platform") {
 				return emptyImage(platform)
@@ -209,10 +208,14 @@ func newV1Image(keychain authn.Keychain, repoName string, platform imgutil.Platf
 		}
 		break
 	}
-
 	return image, nil
 }
 
+func is40x(transportErr *transport.Error) bool {
+	return transportErr.StatusCode == http.StatusNotFound ||
+		transportErr.StatusCode == http.StatusUnauthorized
+}
+
 func emptyImage(platform imgutil.Platform) (v1.Image, error) {
 	cfg := &v1.ConfigFile{
 		Architecture: platform.Architecture,

remote/remote_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"math/rand"
+	"net/http"
 	"os"
 	"testing"
 	"time"
@@ -42,7 +43,6 @@ func TestRemote(t *testing.T) {
 	dockerRegistry = h.NewDockerRegistry(h.WithAuth(dockerConfigDir), sharedStorageOpt)
 	dockerRegistry.Start(t)
 	defer dockerRegistry.Stop(t)
-
 	readonlyDockerRegistry = h.NewDockerRegistry(sharedStorageOpt)
 	readonlyDockerRegistry.Start(t)
 	defer readonlyDockerRegistry.Stop(t)
@@ -501,6 +501,40 @@ func testImage(t *testing.T, when spec.G, it spec.S) {
 				})
 			})
 		})
+
+		when("#RetryLogic", func() {
+			var mockServer *h.MockServer
+
+			when("Manifest API in registry returns status code 200", func() {
+				it.Before(func() {
+					mockServer, repoName = h.SetUpMockServer(t, "org/do-not-retry", http.StatusOK, 0)
+				})
+
+				it("Do not retry after status code 200", func() {
+					VerifyRetryLogic(t, mockServer, repoName, 1)
+				})
+			})
+
+			when("Manifest API in registry returns status code 404", func() {
+				it.Before(func() {
+					mockServer, repoName = h.SetUpMockServer(t, "org/retry-not-found", http.StatusNotFound, 2)
+				})
+
+				it("Retry after status code 404", func() {
+					VerifyRetryLogic(t, mockServer, repoName, 3)
+				})
+			})
+
+			when("Manifest API in registry returns status code 401", func() {
+				it.Before(func() {
+					mockServer, repoName = h.SetUpMockServer(t, "org/retry-unauthorized", http.StatusUnauthorized, 2)
+				})
+
+				it("Retry after status code 401", func() {
+					VerifyRetryLogic(t, mockServer, repoName, 3)
+				})
+			})
+		})
 	})
 
 	when("#Entrypoint", func() {
@@ -1496,3 +1530,11 @@ func testImage(t *testing.T, when spec.G, it spec.S) {
 		})
 	})
 }
+
+func VerifyRetryLogic(t *testing.T, mockServer *h.MockServer, repoName string, expectedCount int) {
+	defer mockServer.Server().Close()
+	_, err := remote.NewImage(repoName, authn.DefaultKeychain, remote.WithPreviousImage(repoName))
+
+	h.AssertNil(t, err)
+	h.AssertEq(t, mockServer.ActualCount(), expectedCount)
+}

testhelpers/mock_registry.go

@@ -0,0 +1,55 @@
+package testhelpers
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+
+	"github.com/google/go-containerregistry/pkg/v1/random"
+)
+
+type MockServer struct {
+	repo        string
+	statusCode  int
+	failedCount int
+	actualCount int
+	server      *httptest.Server
+}
+
+func NewMockServer(repo string, statusCode, failedCount int) *MockServer {
+	return &MockServer{
+		repo:        repo,
+		statusCode:  statusCode,
+		failedCount: failedCount,
+		actualCount: 0,
+	}
+}
+
+func (m *MockServer) Init() *httptest.Server {
+	manifestPath := fmt.Sprintf("/v2/%s/manifests/latest", m.repo)
+	img, _ := random.Image(1024, 1)
+	m.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var err error
+		if r.URL.Path == manifestPath {
+			m.actualCount++
+			if m.actualCount <= m.failedCount {
+				w.WriteHeader(m.statusCode)
+			} else {
+				mm, _ := img.RawManifest()
+				_, err = w.Write(mm)
+			}
+		}
+		if err != nil {
+			fmt.Printf("There was an error in the mock registry %s\n", err)
+		}
+	}))
+	return m.server
+}
+
+func (m *MockServer) ActualCount() int {
+	return m.actualCount
+}
+
+func (m *MockServer) Server() *httptest.Server {
+	return m.server
+}

testhelpers/testhelpers.go

@@ -13,6 +13,7 @@ import (
 	"io/ioutil"
 	"math/rand"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 	"strings"
@@ -409,3 +410,14 @@ func checkResponseError(r io.Reader) error {
 
 	return nil
 }
+
+func SetUpMockServer(t *testing.T, repo string, statusCode, failedCount int) (*MockServer, string) {
+	mockServer := NewMockServer(repo, statusCode, failedCount)
+	server := mockServer.Init()
+	u, err := url.Parse(server.URL)
+
+	AssertNil(t, err)
+
+	repoName := u.Hostname() + ":" + u.Port() + "/" + repo
+	return mockServer, repoName
+}