Add exclued node flag for load image command

Signed-off-by: Monokaix <[email protected]>

pkg/cluster/nodeutils/util.go

@@ -154,3 +154,14 @@ 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()
 }
+
+// RemoveNode remove a node from a node list.
+func RemoveNode(nodeName string, nodes []nodes.Node) []nodes.Node {
+	for i := 0; i < len(nodes); i++ {
+		if nodes[i].String() == nodeName {
+			nodes = append(nodes[:i], nodes[i+1:]...)
+			i--
+		}
+	}
+	return nodes
+}

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

@@ -43,8 +43,9 @@ type (
 )
 
 type flagpole struct {
-	Name  string
-	Nodes []string
+	Name         string
+	Nodes        []string
+	ExcludeNodes []string
 }
 
 // NewCommand returns a new cobra.Command for loading an image into a cluster
@@ -78,6 +79,12 @@ func NewCommand(logger log.Logger, streams cmd.IOStreams) *cobra.Command {
 		nil,
 		"comma separated list of nodes to load images into",
 	)
+	cmd.Flags().StringSliceVar(
+		&flags.ExcludeNodes,
+		"exclude-nodes",
+		nil,
+		"comma separated list of exclude nodes not to load images into",
+	)
 	return cmd
 }
 
@@ -129,6 +136,16 @@ func runE(logger log.Logger, flags *flagpole, args []string) error {
 		}
 	}
 
+	if len(flags.ExcludeNodes) > 0 {
+		for _, name := range flags.ExcludeNodes {
+			_, ok := nodesByName[name]
+			if !ok {
+				return fmt.Errorf("unknown node: %q", name)
+			}
+			candidateNodes = nodeutils.RemoveNode(name, candidateNodes)
+		}
+	}
+
 	// pick only the nodes that don't have the image
 	selectedNodes := map[string]nodes.Node{}
 	fns := []func() error{}

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

@@ -35,8 +35,9 @@ import (
 )
 
 type flagpole struct {
-	Name  string
-	Nodes []string
+	Name         string
+	Nodes        []string
+	ExcludeNodes []string
 }
 
 // NewCommand returns a new cobra.Command for loading an image into a cluster
@@ -75,6 +76,12 @@ func NewCommand(logger log.Logger, streams cmd.IOStreams) *cobra.Command {
 		nil,
 		"comma separated list of nodes to load images into",
 	)
+	cmd.Flags().StringSliceVar(
+		&flags.ExcludeNodes,
+		"exclude-nodes",
+		nil,
+		"comma separated list of exclude nodes not to load images into",
+	)
 	return cmd
 }
 
@@ -91,14 +98,14 @@ func runE(logger log.Logger, flags *flagpole, args []string) error {
 	}
 
 	for _, imageTarPath := range args {
-		if err := loadArchiveToNodes(logger, provider, flags.Name, flags.Nodes, imageTarPath); err != nil {
+		if err := loadArchiveToNodes(logger, provider, flags.Name, flags.Nodes, flags.ExcludeNodes, imageTarPath); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func loadArchiveToNodes(logger log.Logger, provider *cluster.Provider, clusterName string, nodeNames []string, imageArchivePath string) error {
+func loadArchiveToNodes(logger log.Logger, provider *cluster.Provider, clusterName string, nodeNames []string, excludeNodeNames []string, imageArchivePath string) error {
 	// Check if the cluster nodes exist
 	nodeList, err := provider.ListInternalNodes(clusterName)
 	if err != nil {
@@ -130,6 +137,16 @@ func loadArchiveToNodes(logger log.Logger, provider *cluster.Provider, clusterNa
 		}
 	}
 
+	if len(excludeNodeNames) > 0 {
+		for _, name := range excludeNodeNames {
+			_, ok := nodesByName[name]
+			if !ok {
+				return fmt.Errorf("unknown node: %q", name)
+			}
+			selectedNodes = nodeutils.RemoveNode(name, selectedNodes)
+		}
+	}
+
 	// Load the image on the selected nodes
 	fns := []func() error{}
 	for _, selectedNode := range selectedNodes {