GIT-2663: optimized image load with re-tag support

pkg/cluster/nodeutils/util.go

@@ -126,3 +126,31 @@ func ImageID(n nodes.Node, image string) (string, error) {
 	}
 	return crictlOut.Status.ID, nil
 }
+
+// ImageTags is used to perform a reverse lookup of the ImageID to list set of available
+// RepoTags corresponding to the ImageID in question
+func ImageTags(n nodes.Node, imageID string) (map[string]bool, error) {
+	var out bytes.Buffer
+	tags := make(map[string]bool, 0)
+	if err := n.Command("crictl", "inspecti", imageID).SetStdout(&out).Run(); err != nil {
+		return tags, err
+	}
+	crictlOut := struct {
+		Status struct {
+			RepoTags []string `json:"repoTags"`
+		} `json:"status"`
+	}{}
+	if err := json.Unmarshal(out.Bytes(), &crictlOut); err != nil {
+		return tags, err
+	}
+	for _, tag := range crictlOut.Status.RepoTags {
+		tags[tag] = true
+	}
+	return tags, nil
+}
+
+// ReTagImage is used to tag an ImageID with a custom tag specified by imageName parameter
+func ReTagImage(n nodes.Node, imageID, imageName string) error {
+	var out bytes.Buffer
+	return n.Command("ctr", "--namespace=k8s.io", "images", "tag", "--force", imageID, imageName).SetStdout(&out).Run()
+}

pkg/cmd/kind/load/docker-image/docker-image.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/spf13/cobra"
 
@@ -37,6 +38,10 @@ import (
 	"sigs.k8s.io/kind/pkg/internal/runtime"
 )
 
+type (
+	imageTagFetcher func(nodes.Node, string) (map[string]bool, error)
+)
+
 type flagpole struct {
 	Name  string
 	Nodes []string
@@ -128,16 +133,34 @@ func runE(logger log.Logger, flags *flagpole, args []string) error {
 	fns := []func() error{}
 	for i, imageName := range imageNames {
 		imageID := imageIDs[i]
+		processed := false
 		for _, node := range candidateNodes {
+			exists, reTagRequired := checkIfImageReTagRequired(node, imageID, imageName, nodeutils.ImageTags)
+			if exists && !reTagRequired {
+				continue
+			}
+
+			if reTagRequired {
+				// We will try to re-tag the image. If the re-tag fails, we will fall back to the default behavior of loading
+				// the images into the nodes again
+				logger.V(0).Infof("Image with ID: %s already present on the node %s but is missing the tag %s. re-tagging...", imageID, node.String(), imageName)
+				if err := nodeutils.ReTagImage(node, imageID, imageName); err != nil {
+					logger.Errorf("failed to re-tag image on the node %s due to an error %s. Will load it instead...", node.String(), err)
+					selectedNodes = append(selectedNodes, node)
+				} else {
+					processed = true
+				}
+				continue
+			}
 			id, err := nodeutils.ImageID(node, imageName)
 			if err != nil || id != imageID {
 				selectedNodes = append(selectedNodes, node)
 				logger.V(0).Infof("Image: %q with ID %q not yet present on node %q, loading...", imageName, imageID, node.String())
 			}
+			continue
 		}
-		if len(selectedNodes) == 0 {
+		if len(selectedNodes) == 0 && !processed {
 			logger.V(0).Infof("Image: %q with ID %q found to be already present on all nodes.", imageName, imageID)
-			continue
 		}
 	}
 
@@ -215,3 +238,46 @@ func removeDuplicates(slice []string) []string {
 	}
 	return result
 }
+
+// checkIfImageExists makes sure we only perform the reverse lookup of the ImageID to tag map
+func checkIfImageReTagRequired(node nodes.Node, imageID, imageName string, tagFetcher imageTagFetcher) (exists, reTagRequired bool) {
+	tags, err := tagFetcher(node, imageID)
+	if len(tags) == 0 || err != nil {
+		exists = false
+		return
+	}
+	exists = true
+	imageName = sanitizeImage(imageName)
+	if ok := tags[imageName]; ok {
+		reTagRequired = false
+		return
+	}
+	reTagRequired = true
+	return
+}
+
+// sanitizeImage is a helper to return human readable image name
+// This is a modified version of the same function found under providers/podman/images.go
+func sanitizeImage(image string) (sanitizedName string) {
+	const (
+		defaultDomain    = "docker.io/"
+		officialRepoName = "library"
+	)
+	sanitizedName = image
+
+	if !strings.ContainsRune(image, '/') {
+		sanitizedName = officialRepoName + "/" + image
+	}
+
+	i := strings.IndexRune(sanitizedName, '/')
+	if i == -1 || (!strings.ContainsAny(sanitizedName[:i], ".:") && sanitizedName[:i] != "localhost") {
+		sanitizedName = defaultDomain + sanitizedName
+	}
+
+	i = strings.IndexRune(sanitizedName, ':')
+	if i == -1 {
+		sanitizedName += ":latest"
+	}
+
+	return
+}

pkg/cmd/kind/load/docker-image/docker-image_test.go

@@ -17,9 +17,12 @@ limitations under the License.
 package load
 
 import (
+	"errors"
 	"reflect"
 	"sort"
 	"testing"
+
+	"sigs.k8s.io/kind/pkg/cluster/nodes"
 )
 
 func Test_removeDuplicates(t *testing.T) {
@@ -60,3 +63,118 @@ func Test_removeDuplicates(t *testing.T) {
 		})
 	}
 }
+
+func Test_sanitizeImage(t *testing.T) {
+	tests := []struct {
+		name           string
+		image          string
+		sanitizedImage string
+	}{
+		{
+			image:          "ubuntu:18.04",
+			sanitizedImage: "docker.io/library/ubuntu:18.04",
+		},
+		{
+			image:          "custom/ubuntu:18.04",
+			sanitizedImage: "docker.io/custom/ubuntu:18.04",
+		},
+		{
+			image:          "registry.k8s.io/kindest/node:latest",
+			sanitizedImage: "registry.k8s.io/kindest/node:latest",
+		},
+		{
+			image:          "k8s.gcr.io/pause:3.6",
+			sanitizedImage: "k8s.gcr.io/pause:3.6",
+		},
+		{
+			image:          "baz",
+			sanitizedImage: "docker.io/library/baz:latest",
+		},
+		{
+			image:          "other-registry/baz",
+			sanitizedImage: "docker.io/other-registry/baz:latest",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := sanitizeImage(tt.image)
+			if got != tt.sanitizedImage {
+				t.Errorf("sanitizeImage(%s) = %s, want %s", tt.image, got, tt.sanitizedImage)
+			}
+		})
+	}
+}
+
+func Test_checkIfImageReTagRequired(t *testing.T) {
+	tests := []struct {
+		name      string
+		imageTags struct {
+			tags map[string]bool
+			err  error
+		}
+		imageID      string
+		imageName    string
+		returnValues []bool
+	}{
+		{
+			name: "image is already present",
+			imageTags: struct {
+				tags map[string]bool
+				err  error
+			}{
+				map[string]bool{
+					"docker.io/library/image1:tag1": true,
+					"k8s.io/image1:tag1":            true,
+				},
+				nil,
+			},
+			imageID:      "sha256:fd3fd9ab134a864eeb7b2c073c0d90192546f597c60416b81fc4166cca47f29a",
+			imageName:    "k8s.io/image1:tag1",
+			returnValues: []bool{true, false},
+		},
+		{
+			name: "re-tag is required",
+			imageTags: struct {
+				tags map[string]bool
+				err  error
+			}{
+				map[string]bool{
+					"docker.io/library/image1:tag1": true,
+					"k8s.io/image1:tag1":            true,
+				},
+				nil,
+			},
+			imageID:      "sha256:fd3fd9ab134a864eeb7b2c073c0d90192546f597c60416b81fc4166cca47f29a",
+			imageName:    "k8s.io/image1:tag2",
+			returnValues: []bool{true, true},
+		},
+		{
+			name: "image tag fetch failed",
+			imageTags: struct {
+				tags map[string]bool
+				err  error
+			}{
+				map[string]bool{},
+				errors.New("some runtime error"),
+			},
+			imageID:      "sha256:fd3fd9ab134a864eeb7b2c073c0d90192546f597c60416b81fc4166cca47f29a",
+			imageName:    "k8s.io/image1:tag2",
+			returnValues: []bool{false, false},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			// checkIfImageReTagRequired doesn't use the `nodes.Node` type for anything. So
+			// passing a nil value here should be fine as the other two functions that use the
+			// nodes.Node has been stubbed out already
+			exists, reTagRequired := checkIfImageReTagRequired(nil, tc.imageID, tc.imageName, func(n nodes.Node, s string) (map[string]bool, error) {
+				return tc.imageTags.tags, tc.imageTags.err
+			})
+			if exists != tc.returnValues[0] || reTagRequired != tc.returnValues[1] {
+				t.Errorf("checkIfImageReTagRequired failed. Expected: %v, got: [%v, %v]", tc.returnValues, exists, reTagRequired)
+			}
+		})
+	}
+}