Merge pull request #2587 from etungsten/check-before-pull

During kubelet prestart, skip pause image pull if image exist
bottlerocket-os · Dec 2, 2022 · 7050b81 · 7050b81
2 parents 8c92057 + 4ca408f
commit 7050b81
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 19 deletions.
diff --git a/packages/kubernetes-1.21/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.21/prestart-pull-pause-ctr.conf
@@ -5,4 +5,5 @@ ExecStartPre=/usr/bin/host-ctr \
     --namespace=k8s.io \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
-    --registry-config=/etc/host-containers/host-ctr.toml
+    --registry-config=/etc/host-containers/host-ctr.toml \
+    --skip-if-image-exists=true
diff --git a/packages/kubernetes-1.22/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.22/prestart-pull-pause-ctr.conf
@@ -5,4 +5,5 @@ ExecStartPre=/usr/bin/host-ctr \
     --namespace=k8s.io \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
-    --registry-config=/etc/host-containers/host-ctr.toml
+    --registry-config=/etc/host-containers/host-ctr.toml \
+    --skip-if-image-exists=true
diff --git a/packages/kubernetes-1.23/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.23/prestart-pull-pause-ctr.conf
@@ -5,4 +5,5 @@ ExecStartPre=/usr/bin/host-ctr \
     --namespace=k8s.io \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
-    --registry-config=/etc/host-containers/host-ctr.toml
+    --registry-config=/etc/host-containers/host-ctr.toml \
+    --skip-if-image-exists=true
diff --git a/packages/kubernetes-1.24/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.24/prestart-pull-pause-ctr.conf
@@ -5,4 +5,5 @@ ExecStartPre=/usr/bin/host-ctr \
     --namespace=k8s.io \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
-    --registry-config=/etc/host-containers/host-ctr.toml
+    --registry-config=/etc/host-containers/host-ctr.toml \
+    --skip-if-image-exists=true
diff --git a/sources/host-ctr/cmd/host-ctr/main.go b/sources/host-ctr/cmd/host-ctr/main.go
@@ -68,6 +68,7 @@ func App() *cli.App {
 		superpowered     bool
 		registryConfig   string
 		cType            string
+		useCachedImage   bool
 	)
 
 	app := cli.NewApp()
@@ -127,9 +128,15 @@ func App() *cli.App {
 					Destination: &cType,
 					Value:       "host",
 				},
+				&cli.BoolFlag{
+					Name:        "use-cached-image",
+					Usage:       "skips registry authentication and image pull if the image already exists in the image store",
+					Destination: &useCachedImage,
+					Value:       false,
+				},
 			},
 			Action: func(c *cli.Context) error {
-				return runCtr(containerdSocket, namespace, containerID, source, superpowered, registryConfig, containerType(cType))
+				return runCtr(containerdSocket, namespace, containerID, source, superpowered, registryConfig, containerType(cType), useCachedImage)
 			},
 		},
 		{
@@ -148,9 +155,15 @@ func App() *cli.App {
 					Usage:       "path to image registry configuration",
 					Destination: &registryConfig,
 				},
+				&cli.BoolFlag{
+					Name:        "skip-if-image-exists",
+					Usage:       "skips registry authentication and image pull if the image already exists in the image store",
+					Destination: &useCachedImage,
+					Value:       false,
+				},
 			},
 			Action: func(c *cli.Context) error {
-				return pullImageOnly(containerdSocket, namespace, source, registryConfig)
+				return pullImageOnly(containerdSocket, namespace, source, registryConfig, useCachedImage)
 			},
 		},
 		{
@@ -215,7 +228,7 @@ func (ct containerType) Prefix() string {
 	return ""
 }
 
-func runCtr(containerdSocket string, namespace string, containerID string, source string, superpowered bool, registryConfigPath string, cType containerType) error {
+func runCtr(containerdSocket string, namespace string, containerID string, source string, superpowered bool, registryConfigPath string, cType containerType, useCachedImage bool) error {
 	// Check if the containerType provided is valid
 	if !cType.IsValid() {
 		return errors.New("Invalid container type")
@@ -248,16 +261,16 @@ func runCtr(containerdSocket string, namespace string, containerID string, sourc
 	}
 	defer client.Close()
 
-	// Parse the source ref if it looks like an ECR ref.
+	// Check if the image source is an ECR image. If it is, then we need to handle it with the ECR resolver.
 	isECRImage := ecrRegex.MatchString(source)
 	var img containerd.Image
 	if isECRImage {
-		img, err = pullECRImage(ctx, source, client, registryConfigPath)
+		img, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage)
 		if err != nil {
 			return err
 		}
 	} else {
-		img, err = pullImage(ctx, source, client, registryConfigPath)
+		img, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage)
 		if err != nil {
 			log.G(ctx).WithField("ref", source).Error(err)
 			return err
@@ -482,7 +495,7 @@ func runCtr(containerdSocket string, namespace string, containerID string, sourc
 }
 
 // pullImageOnly pulls the specified container image
-func pullImageOnly(containerdSocket string, namespace string, source string, registryConfigPath string) error {
+func pullImageOnly(containerdSocket string, namespace string, source string, registryConfigPath string, useCachedImage bool) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	ctx = namespaces.WithNamespace(ctx, namespace)
@@ -493,18 +506,17 @@ func pullImageOnly(containerdSocket string, namespace string, source string, reg
 	}
 	defer client.Close()
 
-	// Parse the source ref if it looks like an ECR ref.
+	// Check if the image source is an ECR image. If it is, then we need to handle it with the ECR resolver.
 	isECRImage := ecrRegex.MatchString(source)
-	ref := source
 	if isECRImage {
-		_, err = pullECRImage(ctx, source, client, registryConfigPath)
+		_, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage)
 		if err != nil {
 			return err
 		}
 	} else {
-		_, err = pullImage(ctx, ref, client, registryConfigPath)
+		_, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage)
 		if err != nil {
-			log.G(ctx).WithField("ref", ref).Error(err)
+			log.G(ctx).WithField("ref", source).Error(err)
 			return err
 		}
 	}
@@ -564,8 +576,8 @@ func cleanUp(containerdSocket string, namespace string, containerID string) erro
 	return nil
 }
 
-// pullECRImage does some additional conversions before resolving the image reference and pulls the image.
-func pullECRImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string) (containerd.Image, error) {
+// fetchECRImage does some additional conversions before resolving the image reference and fetches the image.
+func fetchECRImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, fetchCachedImageIfExist bool) (containerd.Image, error) {
 	ref := source
 	ecrRef, err := ecr.ParseImageURI(ref)
 	if err != nil {
@@ -579,7 +591,7 @@ func pullECRImage(ctx context.Context, source string, client *containerd.Client,
 		WithField("source", source).
 		Debug("parsed ECR reference from URI")
 
-	img, err := pullImage(ctx, ref, client, registryConfigPath)
+	img, err := fetchImage(ctx, ref, client, registryConfigPath, fetchCachedImageIfExist)
 	if err != nil {
 		log.G(ctx).WithField("ref", ref).Error(err)
 		return nil, err
@@ -882,6 +894,25 @@ func withProxyEnv() oci.SpecOpts {
 	return oci.Compose(withHttpsProxy, withNoProxy)
 }
 
+// fetchImage returns a `containerd.Image` given an image source.
+func fetchImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, useCachedImage bool) (containerd.Image, error) {
+	// Check the containerd image store to see if image exists
+	img, err := client.GetImage(ctx, source)
+	if err != nil {
+		if errdefs.IsNotFound(err) {
+			log.G(ctx).WithField("ref", source).Info("Image does not exist, proceeding to pull image from source.")
+		} else {
+			log.G(ctx).WithField("ref", source).Error(err)
+			return nil, err
+		}
+	}
+	if img != nil && useCachedImage {
+		log.G(ctx).WithField("ref", source).Info("Image exists, fetching cached image from image store")
+		return img, err
+	}
+	return pullImage(ctx, source, client, registryConfigPath)
+}
+
 // pullImage pulls an image from the specified source.
 func pullImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string) (containerd.Image, error) {
 	// Handle registry config