impl: add a retry with result function (#2837)

README.md

@@ -112,7 +112,8 @@ _If you are interested in contributing to kaniko, see
       - [Flag `--ignore-var-run`](#flag---ignore-var-run)
       - [Flag `--ignore-path`](#flag---ignore-path)
       - [Flag `--image-fs-extract-retry`](#flag---image-fs-extract-retry)
-    - [Debug Image](#debug-image)
+      - [Flag `--image-download-retry`](#flag---image-download-retry)    
+  - [Debug Image](#debug-image)
   - [Security](#security)
     - [Verifying Signed Kaniko Images](#verifying-signed-kaniko-images)
   - [Kaniko Builds - Profiling](#kaniko-builds---profiling)
@@ -1094,6 +1095,10 @@ snapshot. Set it multiple times for multiple ignore paths.
 Set this flag to the number of retries that should happen for the extracting an
 image filesystem. Defaults to `0`.
 
+#### Flag `--image-download-retry`
+
+Set this flag to the number of retries that should happen when downloading the remote image. Defaults to `0`.
+
 ### Debug Image
 
 The kaniko executor image is based on scratch and doesn't contain a shell. We

cmd/executor/cmd/root.go

@@ -212,6 +212,7 @@ func addKanikoOptionsFlags() {
 	RootCmd.PersistentFlags().BoolVarP(&opts.SkipTLSVerifyPull, "skip-tls-verify-pull", "", false, "Pull from insecure registry ignoring TLS verify")
 	RootCmd.PersistentFlags().IntVar(&opts.PushRetry, "push-retry", 0, "Number of retries for the push operation")
 	RootCmd.PersistentFlags().IntVar(&opts.ImageFSExtractRetry, "image-fs-extract-retry", 0, "Number of retries for image FS extraction")
+	RootCmd.PersistentFlags().IntVar(&opts.ImageDownloadRetry, "image-download-retry", 0, "Number of retries for downloading the remote image")
 	RootCmd.PersistentFlags().StringVarP(&opts.KanikoDir, "kaniko-dir", "", constants.DefaultKanikoPath, "Path to the kaniko directory, this takes precedence over the KANIKO_DIR environment variable.")
 	RootCmd.PersistentFlags().StringVarP(&opts.TarPath, "tar-path", "", "", "Path to save the image in as a tarball instead of pushing")
 	RootCmd.PersistentFlags().BoolVarP(&opts.SingleSnapshot, "single-snapshot", "", false, "Take a single snapshot at the end of the build.")

pkg/config/options.go

@@ -43,6 +43,7 @@ type RegistryOptions struct {
 	InsecurePull                 bool
 	SkipTLSVerifyPull            bool
 	PushRetry                    int
+	ImageDownloadRetry           int
 }
 
 // KanikoOptions are options that are set by command line arguments

pkg/image/remote/remote.go

@@ -70,8 +70,13 @@ func RetrieveRemoteImage(image string, opts config.RegistryOptions, customPlatfo
 			ref := setNewRegistry(ref, newReg)
 
 			logrus.Infof("Retrieving image %s from registry mirror %s", ref, registryMirror)
-			remoteImage, err := remoteImageFunc(ref, remoteOptions(registryMirror, opts, customPlatform)...)
-			if err != nil {
+			retryFunc := func() (v1.Image, error) {
+				return remoteImageFunc(ref, remoteOptions(registryMirror, opts, customPlatform)...)
+			}
+
+			var remoteImage v1.Image
+			var err error
+			if remoteImage, err = util.RetryWithResult(retryFunc, opts.ImageDownloadRetry, 1000); err != nil {
 				logrus.Warnf("Failed to retrieve image %s from registry mirror %s: %s. Will try with the next mirror, or fallback to the default registry.", ref, registryMirror, err)
 				continue
 			}
@@ -97,9 +102,12 @@ func RetrieveRemoteImage(image string, opts config.RegistryOptions, customPlatfo
 
 	logrus.Infof("Retrieving image %s from registry %s", ref, registryName)
 
-	remoteImage, err := remoteImageFunc(ref, remoteOptions(registryName, opts, customPlatform)...)
+	retryFunc := func() (v1.Image, error) {
+		return remoteImageFunc(ref, remoteOptions(registryName, opts, customPlatform)...)
+	}
 
-	if remoteImage != nil {
+	var remoteImage v1.Image
+	if remoteImage, err = util.RetryWithResult(retryFunc, opts.ImageDownloadRetry, 1000); remoteImage != nil {
 		manifestCache[image] = remoteImage
 	}
 

pkg/image/remote/remote_test.go

@@ -27,8 +27,11 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
+const image string = "debian"
+
 // mockImage mocks the v1.Image interface
-type mockImage struct{}
+type mockImage struct {
+}
 
 func (m *mockImage) ConfigFile() (*v1.ConfigFile, error) {
 	return nil, nil
@@ -79,7 +82,6 @@ func (m *mockImage) Size() (int64, error) {
 }
 
 func Test_normalizeReference(t *testing.T) {
-	image := "debian"
 	expected := "index.docker.io/library/debian:latest"
 
 	ref, err := name.ParseReference(image)
@@ -112,7 +114,6 @@ func Test_RetrieveRemoteImage_manifestCache(t *testing.T) {
 }
 
 func Test_RetrieveRemoteImage_skipFallback(t *testing.T) {
-	image := "debian"
 	registryMirror := "some-registry"
 
 	opts := config.RegistryOptions{
@@ -140,3 +141,45 @@ func Test_RetrieveRemoteImage_skipFallback(t *testing.T) {
 		t.Fatal("Expected call to fail because fallback to default registry is skipped")
 	}
 }
+
+func Test_RetryRetrieveRemoteImageSucceeds(t *testing.T) {
+	opts := config.RegistryOptions{
+		ImageDownloadRetry: 2,
+	}
+	attempts := 0
+	remoteImageFunc = func(ref name.Reference, options ...remote.Option) (v1.Image, error) {
+		if attempts < 2 {
+			attempts++
+			return nil, errors.New("no image found")
+		}
+		return &mockImage{}, nil
+	}
+
+	// Clean cached image
+	manifestCache = make(map[string]v1.Image)
+
+	if _, err := RetrieveRemoteImage(image, opts, ""); err != nil {
+		t.Fatal("Expected call to succeed because of retry")
+	}
+}
+
+func Test_NoRetryRetrieveRemoteImageFails(t *testing.T) {
+	opts := config.RegistryOptions{
+		ImageDownloadRetry: 0,
+	}
+	attempts := 0
+	remoteImageFunc = func(ref name.Reference, options ...remote.Option) (v1.Image, error) {
+		if attempts < 1 {
+			attempts++
+			return nil, errors.New("no image found")
+		}
+		return &mockImage{}, nil
+	}
+
+	// Clean cached image
+	manifestCache = make(map[string]v1.Image)
+
+	if _, err := RetrieveRemoteImage(image, opts, ""); err == nil {
+		t.Fatal("Expected call to fail because there is no retry")
+	}
+}