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

GIT-2663: sanitizeImage function to help generate fully qualified image names

GIT-2663: sanitizeImage function to help generate fully qualified image names

GIT-2663: add comments for sanitizeImage function and add UTs

GIT-2663: add comments for sanitizeImage function and add UTs

pkg/cluster/nodeutils/util.go

@@ -126,3 +126,41 @@ 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
+}
+
+// ImageReTagSupported checks if the version of the ctr running in the cluster has support
+// for performing ctr images tag command
+func ImageReTagSupported(n nodes.Node) bool {
+	var out bytes.Buffer
+	if err := n.Command("ctr", "images", "tag", "--help").SetStdout(&out).Run(); err != nil {
+		return false
+	}
+	return true
+}
+
+// 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,11 @@ import (
 	"sigs.k8s.io/kind/pkg/internal/runtime"
 )
 
+var (
+	imageReTagSupportChecker = nodeutils.ImageReTagSupported
+	imageTagFetcher          = nodeutils.ImageTags
+)
+
 type flagpole struct {
 	Name  string
 	Nodes []string
@@ -128,16 +134,33 @@ 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 {
-			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())
+			exists, reTagRequired := checkIfImageReTagRequired(node, imageID, imageName)
+			if !exists {
+				// This could be because of one of two reasons. Either the previous command ran into an error or no image
+				// with matching Image ID was found. Let us try again and see if the image is indeed required to be loaded
+				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 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
+				}
 			}
 		}
-		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,49 @@ func removeDuplicates(slice []string) []string {
 	}
 	return result
 }
+
+// checkIfImageExists makes sure we only perform the reverse lookup of the ImageID to tag map in case if
+// the underlying version of ctr supports image tag workflow.
+func checkIfImageReTagRequired(node nodes.Node, imageID, imageName string) (exists bool, reTagRequired bool) {
+	if !imageReTagSupportChecker(node) {
+		return
+	}
+	tags, err := imageTagFetcher(node, imageID)
+	if len(tags) == 0 || err != nil {
+		return
+	}
+	imageName = sanitizeImage(imageName)
+	if ok := tags[imageName]; ok {
+		exists = true
+		return
+	}
+	exists = true
+	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,138 @@ 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
+		reTagSupported bool
+		imageTags      struct {
+			tags map[string]bool
+			err  error
+		}
+		imageID       string
+		imageName     string
+		exists        bool
+		reTagRequired bool
+	}{
+		{
+			name:           "re-tag is not supported",
+			reTagSupported: false,
+		},
+		{
+			name:           "image is already present",
+			reTagSupported: true,
+			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",
+			exists:        true,
+			reTagRequired: false,
+		},
+		{
+			name:           "re-tag is required",
+			reTagSupported: true,
+			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",
+			exists:        true,
+			reTagRequired: true,
+		},
+		{
+			name:           "image tag fetch failed",
+			reTagSupported: true,
+			imageTags: struct {
+				tags map[string]bool
+				err  error
+			}{
+				map[string]bool{},
+				errors.New("some runtime error"),
+			},
+			imageID:       "sha256:fd3fd9ab134a864eeb7b2c073c0d90192546f597c60416b81fc4166cca47f29a",
+			imageName:     "k8s.io/image1:tag2",
+			exists:        false,
+			reTagRequired: false,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			// Doing the following patching to avoid having to mock teh nodes.Node type for UT
+			imageReTagSupportChecker = func(n nodes.Node) bool {
+				return tc.reTagSupported
+			}
+			imageTagFetcher = func(n nodes.Node, imageID string) (map[string]bool, error) {
+				return tc.imageTags.tags, tc.imageTags.err
+			}
+			// 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)
+			if exists != tc.exists {
+				t.Errorf("checkIfImageReTagRequired failed. exists Expected: %v, got: %v", tc.exists, exists)
+			}
+			if reTagRequired != tc.reTagRequired {
+				t.Errorf("checkIfImageReTagRequired failed. reTagRequired Expected: %v, got: %v", tc.reTagRequired, reTagRequired)
+			}
+		})
+	}
+}