Use the default retry predicate in transport

internal/retry/retry.go

@@ -76,3 +76,19 @@ func Retry(f func() error, p Predicate, backoff wait.Backoff) (err error) {
 	wait.ExponentialBackoff(backoff, condition)
 	return
 }
+
+type contextKey string
+
+var key = contextKey("never")
+
+// Never returns a context that signals something should not be retried.
+// This is a hack and can be used to communicate across package boundaries
+// to avoid retry amplification.
+func Never(ctx context.Context) context.Context {
+	return context.WithValue(ctx, key, true)
+}
+
+// Ever returns true if the context was wrapped by Never.
+func Ever(ctx context.Context) bool {
+	return ctx.Value(key) == nil
+}

pkg/v1/remote/check.go

@@ -35,11 +35,10 @@ func CheckPushPermission(ref name.Reference, kc authn.Keychain, t http.RoundTrip
 	// authorize a push. Figure out how to return early here when we can,
 	// to avoid a roundtrip for spec-compliant registries.
 	w := writer{
-		repo:    ref.Context(),
-		client:  &http.Client{Transport: tr},
-		context: context.Background(),
+		repo:   ref.Context(),
+		client: &http.Client{Transport: tr},
 	}
-	loc, _, err := w.initiateUpload("", "", "")
+	loc, _, err := w.initiateUpload(context.Background(), "", "", "")
 	if loc != "" {
 		// Since we're only initiating the upload to check whether we
 		// can, we should attempt to cancel it, in case initiating

pkg/v1/remote/multi_write.go

@@ -89,7 +89,6 @@ func MultiWrite(m map[name.Reference]Taggable, options ...Option) (rerr error) {
 	w := writer{
 		repo:      repo,
 		client:    &http.Client{Transport: tr},
-		context:   o.context,
 		backoff:   o.retryBackoff,
 		predicate: o.retryPredicate,
 	}

pkg/v1/remote/options.go

@@ -82,6 +82,14 @@ var fastBackoff = Backoff{
 	Steps:    3,
 }
 
+var retryableStatusCodes = []int{
+	http.StatusRequestTimeout,
+	http.StatusInternalServerError,
+	http.StatusBadGateway,
+	http.StatusServiceUnavailable,
+	http.StatusGatewayTimeout,
+}
+
 const (
 	defaultJobs = 4
 
@@ -151,7 +159,7 @@ func makeOptions(target authn.Resource, opts ...Option) (*options, error) {
 		}
 
 		// Wrap the transport in something that can retry network flakes.
-		o.transport = transport.NewRetry(o.transport)
+		o.transport = transport.NewRetry(o.transport, transport.WithRetryPredicate(defaultRetryPredicate), transport.WithRetryStatusCodes(retryableStatusCodes...))
 
 		// Wrap this last to prevent transport.New from double-wrapping.
 		if o.userAgent != "" {

pkg/v1/remote/transport/retry.go

@@ -21,12 +21,12 @@ import (
 	"github.com/google/go-containerregistry/internal/retry"
 )
 
-// Sleep for 0.1, 0.3, 0.9, 2.7 seconds. This should cover networking blips.
+// Sleep for 0.1 then 0.3 seconds. This should cover networking blips.
 var defaultBackoff = retry.Backoff{
 	Duration: 100 * time.Millisecond,
 	Factor:   3.0,
 	Jitter:   0.1,
-	Steps:    5,
+	Steps:    3,
 }
 
 var _ http.RoundTripper = (*retryTransport)(nil)
@@ -36,6 +36,7 @@ type retryTransport struct {
 	inner     http.RoundTripper
 	backoff   retry.Backoff
 	predicate retry.Predicate
+	codes     []int
 }
 
 // Option is a functional option for retryTransport.
@@ -44,6 +45,7 @@ type Option func(*options)
 type options struct {
 	backoff   retry.Backoff
 	predicate retry.Predicate
+	codes     []int
 }
 
 // Backoff is an alias of retry.Backoff to expose this configuration option to consumers of this lib
@@ -63,6 +65,13 @@ func WithRetryPredicate(predicate func(error) bool) Option {
 	}
 }
 
+// WithRetryStatusCodes sets which http response codes will be retried.
+func WithRetryStatusCodes(codes ...int) Option {
+	return func(o *options) {
+		o.codes = codes
+	}
+}
+
 // NewRetry returns a transport that retries errors.
 func NewRetry(inner http.RoundTripper, opts ...Option) http.RoundTripper {
 	o := &options{
@@ -78,12 +87,23 @@ func NewRetry(inner http.RoundTripper, opts ...Option) http.RoundTripper {
 		inner:     inner,
 		backoff:   o.backoff,
 		predicate: o.predicate,
+		codes:     o.codes,
 	}
 }
 
 func (t *retryTransport) RoundTrip(in *http.Request) (out *http.Response, err error) {
 	roundtrip := func() error {
 		out, err = t.inner.RoundTrip(in)
+		if !retry.Ever(in.Context()) {
+			return nil
+		}
+		if out != nil {
+			for _, code := range t.codes {
+				if out.StatusCode == code {
+					return CheckError(out)
+				}
+			}
+		}
 		return err
 	}
 	retry.Retry(roundtrip, t.predicate, t.backoff)

pkg/v1/remote/transport/retry_test.go

@@ -17,8 +17,10 @@ package transport
 import (
 	"context"
 	"errors"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"
 
@@ -27,12 +29,19 @@ import (
 
 type mockTransport struct {
 	errs  []error
+	resps []*http.Response
 	count int
 }
 
 func (t *mockTransport) RoundTrip(in *http.Request) (out *http.Response, err error) {
 	defer func() { t.count++ }()
-	return nil, t.errs[t.count]
+	if t.count < len(t.resps) {
+		out = t.resps[t.count]
+	}
+	if t.count < len(t.errs) {
+		err = t.errs[t.count]
+	}
+	return
 }
 
 type perm struct{}
@@ -51,11 +60,25 @@ func (e temp) Temporary() bool {
 	return true
 }
 
+func resp(code int) *http.Response {
+	return &http.Response{
+		StatusCode: code,
+		Body:       io.NopCloser(strings.NewReader("hi")),
+	}
+}
+
 func TestRetryTransport(t *testing.T) {
 	for _, test := range []struct {
 		errs  []error
+		resps []*http.Response
+		ctx   context.Context
 		count int
 	}{{
+		// Don't retry retry.Never.
+		errs:  []error{temp{}},
+		ctx:   retry.Never(context.Background()),
+		count: 1,
+	}, {
 		// Don't retry permanent.
 		errs:  []error{perm{}},
 		count: 1,
@@ -67,14 +90,36 @@ func TestRetryTransport(t *testing.T) {
 		// Stop at some max.
 		errs:  []error{temp{}, temp{}, temp{}, temp{}, temp{}},
 		count: 3,
+	}, {
+		// Retry http errors.
+		errs: []error{nil, nil, temp{}, temp{}, temp{}},
+		resps: []*http.Response{
+			resp(http.StatusRequestTimeout),
+			resp(http.StatusInternalServerError),
+			nil,
+		},
+		count: 3,
 	}} {
 		mt := mockTransport{
-			errs: test.errs,
+			errs:  test.errs,
+			resps: test.resps,
 		}
 
-		tr := NewRetry(&mt, WithRetryBackoff(retry.Backoff{Steps: 3}), WithRetryPredicate(retry.IsTemporary))
+		tr := NewRetry(&mt,
+			WithRetryBackoff(retry.Backoff{Steps: 3}),
+			WithRetryPredicate(retry.IsTemporary),
+			WithRetryStatusCodes(http.StatusRequestTimeout, http.StatusInternalServerError),
+		)
 
-		tr.RoundTrip(nil)
+		ctx := context.Background()
+		if test.ctx != nil {
+			ctx = test.ctx
+		}
+		req, err := http.NewRequestWithContext(ctx, "GET", "example.com", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		tr.RoundTrip(req)
 		if mt.count != test.count {
 			t.Errorf("wrong count, wanted %d, got %d", test.count, mt.count)
 		}