host-ctr: add label flag to host-ctr pull-image.

Sandbox container image on k8s 1.29 isn't pinned and causea the Sandbox container image would be GC. we add label flag to host-ctr pull-image and pass "io.cri-containerd.pinned=pinned" label to all 1.29 varaints. Then the Sandbox contianer image will be pinned and avoid GC issue.
bottlerocket-os · Feb 5, 2024 · 44f9c90 · 44f9c90
1 parent 2b1a787
commit 44f9c90
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 15 deletions.
diff --git a/packages/kubernetes-1.29/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.29/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --labels="io.cri-containerd.pinned=pinned"
diff --git a/sources/host-ctr/cmd/host-ctr/host-ctr b/sources/host-ctr/cmd/host-ctr/host-ctr
diff --git a/sources/host-ctr/cmd/host-ctr/main.go b/sources/host-ctr/cmd/host-ctr/main.go
@@ -40,7 +40,7 @@ import (
 var ecrRegex = regexp.MustCompile(`(^[a-zA-Z0-9][a-zA-Z0-9-_]*)\.dkr\.ecr\.([a-zA-Z0-9][a-zA-Z0-9-_]*)\.amazonaws\.com(\.cn)?.*`)
 
 func init() {
-	rand.Seed(time.Now().UnixNano())
+	rand.New(rand.NewSource(time.Now().UnixNano()))
 	// Dispatch logging output instead of writing all levels' messages to
 	// stderr.
 	log.L.Logger.SetOutput(io.Discard)
@@ -69,6 +69,7 @@ func App() *cli.App {
 		registryConfig   string
 		cType            string
 		useCachedImage   bool
+		labels           string
 	)
 
 	app := cli.NewApp()
@@ -161,9 +162,15 @@ func App() *cli.App {
 					Destination: &useCachedImage,
 					Value:       false,
 				},
+				&cli.StringFlag{
+					Name:        "labels",
+					Usage:       "add labels to the pulled image. Each label should be in the format of `key=value` and separated by a comma",
+					Destination: &labels,
+					Value:       "",
+				},
 			},
 			Action: func(c *cli.Context) error {
-				return pullImageOnly(containerdSocket, namespace, source, registryConfig, useCachedImage)
+				return pullImageOnly(containerdSocket, namespace, source, registryConfig, useCachedImage, labels)
 			},
 		},
 		{
@@ -274,13 +281,14 @@ func runCtr(containerdSocket string, namespace string, containerID string, sourc
 	// 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 = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage)
+		img, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage, "")
 		if err != nil {
 			return err
 		}
 	} else {
-		img, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage)
+		img, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage, "")
 		if err != nil {
 			log.G(ctx).WithField("ref", source).Error(err)
 			return err
@@ -494,7 +502,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, useCachedImage bool) error {
+func pullImageOnly(containerdSocket string, namespace string, source string, registryConfigPath string, useCachedImage bool, labels string) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	ctx = namespaces.WithNamespace(ctx, namespace)
@@ -508,12 +516,12 @@ func pullImageOnly(containerdSocket string, namespace string, source string, reg
 	// 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)
 	if isECRImage {
-		_, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage)
+		_, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage, labels)
 		if err != nil {
 			return err
 		}
 	} else {
-		_, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage)
+		_, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage, labels)
 		if err != nil {
 			log.G(ctx).WithField("ref", source).Error(err)
 			return err
@@ -646,7 +654,7 @@ func fetchECRRef(ctx context.Context, input string) (ecr.ECRSpec, 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) {
+func fetchECRImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, fetchCachedImageIfExist bool, labels string) (containerd.Image, error) {
 	ecrRef, err := fetchECRRef(ctx, source)
 	if err != nil {
 		return nil, err
@@ -658,7 +666,7 @@ func fetchECRImage(ctx context.Context, source string, client *containerd.Client
 		WithField("source", source).
 		Debug("parsed ECR reference from URI")
 
-	img, err := fetchImage(ctx, ref, client, registryConfigPath, fetchCachedImageIfExist)
+	img, err := fetchImage(ctx, ref, client, registryConfigPath, fetchCachedImageIfExist, labels)
 	if err != nil {
 		log.G(ctx).WithField("ref", ref).Error(err)
 		return nil, err
@@ -971,7 +979,7 @@ func withProxyEnv() oci.SpecOpts {
 }
 
 // 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) {
+func fetchImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, useCachedImage bool, labels string) (containerd.Image, error) {
 	// Check the containerd image store to see if image exists
 	img, err := client.GetImage(ctx, source)
 	if err != nil {
@@ -986,11 +994,11 @@ func fetchImage(ctx context.Context, source string, client *containerd.Client, r
 		log.G(ctx).WithField("ref", source).Info("Image exists, fetching cached image from image store")
 		return img, err
 	}
-	return pullImage(ctx, source, client, registryConfigPath)
+	return pullImage(ctx, source, client, registryConfigPath, labels)
 }
 
 // pullImage pulls an image from the specified source.
-func pullImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string) (containerd.Image, error) {
+func pullImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, labels string) (containerd.Image, error) {
 	// Handle registry config
 	var registryConfig *RegistryConfig
 	if registryConfigPath != "" {
@@ -1017,9 +1025,18 @@ func pullImage(ctx context.Context, source string, client *containerd.Client, re
 	var img containerd.Image
 	for {
 		var err error
-		img, err = client.Pull(ctx, source,
+
+		pullOpts := []containerd.RemoteOpt{
 			withDynamicResolver(ctx, source, registryConfig),
-			containerd.WithSchema1Conversion)
+			containerd.WithSchema1Conversion,
+		}
+
+		if len(labels) != 0 {
+			pullOpts = append(pullOpts, containerd.WithPullLabels(convertLabels(labels)))
+		}
+
+		img, err = client.Pull(ctx, source, pullOpts...)
+
 		if err == nil {
 			log.G(ctx).WithField("img", img.Name()).Info("pulled image successfully")
 			break
@@ -1184,3 +1201,15 @@ func withUnmaskedPaths(unmaskPaths []string) oci.SpecOpts {
 		return nil
 	}
 }
+
+// Convert labels to map[string]string for containerd.WithPullLabels.
+// Labels are in the format of "key=value,key2=value2".
+func convertLabels(labels string) map[string]string {
+	labelsMap := make(map[string]string)
+	labelsList := strings.Split(labels, ",")
+	for _, label := range labelsList {
+		labelKeyValue := strings.Split(label, "=")
+		labelsMap[labelKeyValue[0]] = labelKeyValue[1]
+	}
+	return labelsMap
+}