From eb953aaa3f9b45af755cacc5ef7661f57b83ca54 Mon Sep 17 00:00:00 2001
From: Jason Hall <jasonhall@redhat.com>
Date: Mon, 20 Dec 2021 12:35:45 -0500
Subject: [PATCH] cache image lookups to avoid unnecessary GETs for image
 configs

---
 go.mod                            |  1 +
 pkg/pod/entrypoint_lookup_impl.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/go.mod b/go.mod
index 51c38b1f0b0..4aef02b26ea 100644
--- a/go.mod
+++ b/go.mod
@@ -16,6 +16,7 @@ require (
 	github.com/google/uuid v1.3.0
 	github.com/googleapis/gnostic v0.5.3 // indirect
 	github.com/hashicorp/go-multierror v1.1.1
+	github.com/hashicorp/golang-lru v0.5.4
 	github.com/jenkins-x/go-scm v1.10.10
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/opencontainers/image-spec v1.0.3-0.20211202222133-eacdcc10569b
diff --git a/pkg/pod/entrypoint_lookup_impl.go b/pkg/pod/entrypoint_lookup_impl.go
index f4dc8b9bf11..59133ae0dc6 100644
--- a/pkg/pod/entrypoint_lookup_impl.go
+++ b/pkg/pod/entrypoint_lookup_impl.go
@@ -27,19 +27,28 @@ import (
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
+	lru "github.com/hashicorp/golang-lru"
 	specs "github.com/opencontainers/image-spec/specs-go/v1"
 	"k8s.io/client-go/kubernetes"
 )
 
+const cacheSize = 1024
+
 type entrypointCache struct {
 	kubeclient kubernetes.Interface
+	lru        *lru.Cache // cache of digest->map[string][]string commands
 }
 
 // NewEntrypointCache returns a new entrypoint cache implementation that uses
 // K8s credentials to pull image metadata from a container image registry.
 func NewEntrypointCache(kubeclient kubernetes.Interface) (EntrypointCache, error) {
+	lru, err := lru.New(cacheSize)
+	if err != nil {
+		return nil, err
+	}
 	return &entrypointCache{
 		kubeclient: kubeclient,
+		lru:        lru,
 	}, nil
 }
 
@@ -49,6 +58,13 @@ func NewEntrypointCache(kubeclient kubernetes.Interface) (EntrypointCache, error
 // reference referred to an index, the returned digest will be the index's
 // digest, not any platform-specific image contained by the index.
 func (e *entrypointCache) Get(ctx context.Context, ref name.Reference, namespace, serviceAccountName string) (*imageData, error) {
+	// If image is specified by digest, check the local cache.
+	if digest, ok := ref.(name.Digest); ok {
+		if id, ok := e.lru.Get(digest.String()); ok {
+			return id.(*imageData), nil
+		}
+	}
+
 	// Consult the remote registry, using imagePullSecrets.
 	kc, err := k8schain.New(ctx, e.kubeclient, k8schain.Options{
 		Namespace:          namespace,
@@ -63,6 +79,15 @@ func (e *entrypointCache) Get(ctx context.Context, ref name.Reference, namespace
 	if err != nil {
 		return nil, err
 	}
+
+	// Check the cache for this ref@digest, in case we've seen it before.
+	// This saves looking up each continuent image's commands if we've seen
+	// the multi-platform image before.
+	refByDigest := ref.Context().Digest(desc.Digest.String()).String()
+	if id, ok := e.lru.Get(refByDigest); ok {
+		return id.(*imageData), nil
+	}
+
 	id := &imageData{
 		digest:   desc.Digest,
 		commands: map[string][]string{},
@@ -112,6 +137,10 @@ func (e *entrypointCache) Get(ctx context.Context, ref name.Reference, namespace
 	default:
 		return nil, errors.New("unsupported media type for image reference")
 	}
+
+	// Cache the digest->commands for future lookup.
+	e.lru.Add(refByDigest, id)
+
 	return id, nil
 }