From 0ef53aa1d37e0e2ae9a15f8eef30bed74503246f Mon Sep 17 00:00:00 2001
From: Sascha Schwarze <schwarzs@de.ibm.com>
Date: Thu, 28 Jan 2021 20:10:13 +0100
Subject: [PATCH] Optimize access to image manifests (#1555)

Introduce an in-memory cache for retrieved manifests in remote.go
---
 pkg/image/image_util.go         |  2 +-
 pkg/image/remote/remote.go      | 23 ++++++++++-
 pkg/image/remote/remote_test.go | 68 +++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/pkg/image/image_util.go b/pkg/image/image_util.go
index 269e2b40d..21ec54ab0 100644
--- a/pkg/image/image_util.go
+++ b/pkg/image/image_util.go
@@ -69,7 +69,7 @@ func RetrieveSourceImage(stage config.KanikoStage, opts *config.KanikoOptions) (
 
 	// Finally, check if local caching is enabled
 	// If so, look in the local cache before trying the remote registry
-	if opts.CacheDir != "" {
+	if opts.Cache && opts.CacheDir != "" {
 		cachedImage, err := cachedImage(opts, currentBaseName)
 		if err != nil {
 			switch {
diff --git a/pkg/image/remote/remote.go b/pkg/image/remote/remote.go
index 97adda377..ea40b2733 100644
--- a/pkg/image/remote/remote.go
+++ b/pkg/image/remote/remote.go
@@ -31,9 +31,20 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+var (
+	manifestCache = make(map[string]v1.Image)
+)
+
 // RetrieveRemoteImage retrieves the manifest for the specified image from the specified registry
 func RetrieveRemoteImage(image string, opts config.RegistryOptions, customPlatform string) (v1.Image, error) {
 	logrus.Infof("Retrieving image manifest %s", image)
+
+	cachedRemoteImage := manifestCache[image]
+	if cachedRemoteImage != nil {
+		logrus.Infof("Returning cached image manifest")
+		return cachedRemoteImage, nil
+	}
+
 	ref, err := name.ParseReference(image, name.WeakValidation)
 	if err != nil {
 		return nil, err
@@ -63,6 +74,9 @@ func RetrieveRemoteImage(image string, opts config.RegistryOptions, customPlatfo
 				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
 			}
+
+			manifestCache[image] = remoteImage
+
 			return remoteImage, nil
 		}
 	}
@@ -77,7 +91,14 @@ func RetrieveRemoteImage(image string, opts config.RegistryOptions, customPlatfo
 	}
 
 	logrus.Infof("Retrieving image %s from registry %s", ref, registryName)
-	return remote.Image(ref, remoteOptions(registryName, opts, customPlatform)...)
+
+	remoteImage, err := remote.Image(ref, remoteOptions(registryName, opts, customPlatform)...)
+
+	if remoteImage != nil {
+		manifestCache[image] = remoteImage
+	}
+
+	return remoteImage, err
 }
 
 // normalizeReference adds the library/ prefix to images without it.
diff --git a/pkg/image/remote/remote_test.go b/pkg/image/remote/remote_test.go
index 74a0fb4ac..3879f017d 100644
--- a/pkg/image/remote/remote_test.go
+++ b/pkg/image/remote/remote_test.go
@@ -19,9 +19,63 @@ package remote
 import (
 	"testing"
 
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
+// mockImage mocks the v1.Image interface
+type mockImage struct{}
+
+func (m *mockImage) ConfigFile() (*v1.ConfigFile, error) {
+	return nil, nil
+}
+
+func (m *mockImage) ConfigName() (v1.Hash, error) {
+	return v1.Hash{}, nil
+}
+
+func (m *mockImage) Descriptor() (*v1.Descriptor, error) {
+	return nil, nil
+}
+
+func (m *mockImage) Digest() (v1.Hash, error) {
+	return v1.Hash{}, nil
+}
+
+func (m *mockImage) LayerByDigest(v1.Hash) (v1.Layer, error) {
+	return nil, nil
+}
+
+func (m *mockImage) LayerByDiffID(v1.Hash) (v1.Layer, error) {
+	return nil, nil
+}
+
+func (m *mockImage) Layers() ([]v1.Layer, error) {
+	return nil, nil
+}
+
+func (m *mockImage) Manifest() (*v1.Manifest, error) {
+	return nil, nil
+}
+
+func (m *mockImage) MediaType() (types.MediaType, error) {
+	return "application/vnd.oci.descriptor.v1+json", nil
+}
+
+func (m *mockImage) RawManifest() ([]byte, error) {
+	return nil, nil
+}
+
+func (m *mockImage) RawConfigFile() ([]byte, error) {
+	return nil, nil
+}
+
+func (m *mockImage) Size() (int64, error) {
+	return 0, nil
+}
+
 func Test_normalizeReference(t *testing.T) {
 	image := "debian"
 	expected := "index.docker.io/library/debian:latest"
@@ -40,3 +94,17 @@ func Test_normalizeReference(t *testing.T) {
 		t.Errorf("%s should have been normalized to %s, got %s", ref2.Name(), expected, ref.Name())
 	}
 }
+
+func Test_RetrieveRemoteImage_manifestCache(t *testing.T) {
+	nonExistingImageName := "this_is_a_non_existing_image_reference"
+
+	if _, err := RetrieveRemoteImage(nonExistingImageName, config.RegistryOptions{}, ""); err == nil {
+		t.Fatal("Expected call to fail because there is no manifest for this image.")
+	}
+
+	manifestCache[nonExistingImageName] = &mockImage{}
+
+	if image, err := RetrieveRemoteImage(nonExistingImageName, config.RegistryOptions{}, ""); image == nil || err != nil {
+		t.Fatal("Expected call to succeed because there is a manifest for this image in the cache.")
+	}
+}