Enable to pass additional handler on pull for stargz-based remote snapshots

Throughout container lifecycle, pulling image is one of the time-consuming
steps. Recently, containerd community started to tackle this issue with
stargz-based remote snapshots, as a non-core
subproject(https://github.com/containerd/stargz-snapshotter).

This snapshotter is implemented as a standard proxy plugin but it requires the
client to pass some additional information (image ref and layer digest) for each
pull operation to query layer contents on the registry. Stargz snapshotter
project provides an image handler to do this and stargz snapshot users need to
pass this handler to containerd client.

This commit enables to use stargz-based remote snapshots through CRI by passing
the handler to containerd client on pull operation.

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>

pkg/config/config.go

@@ -73,6 +73,11 @@ type ContainerdConfig struct {
 	// NoPivot disables pivot-root (linux only), required when running a container in a RamDisk with runc
 	// This only works for runtime type "io.containerd.runtime.v1.linux".
 	NoPivot bool `toml:"no_pivot" json:"noPivot"`
+
+	// DisableSnapshotAnnotations disables to pass additional annotations (image
+	// related information) to snapshotters. These annotations are required by
+	// stargz snapshotter (https://github.com/containerd/stargz-snapshotter).
+	DisableSnapshotAnnotations bool `toml:"disable_snapshot_annotations" json:"disableSnapshotAnnotations"`
 }
 
 // CniConfig contains toml config related to cni

pkg/server/image_pull.go

@@ -120,6 +120,10 @@ func (c *criService) PullImage(ctx context.Context, r *runtime.PullImageRequest)
 	}
 
 	pullOpts = append(pullOpts, c.encryptedImagesPullOpts()...)
+	if !c.config.ContainerdConfig.DisableSnapshotAnnotations {
+		pullOpts = append(pullOpts,
+			containerd.WithImageHandlerWrapper(appendInfoHandlerWrapper(ref)))
+	}
 
 	image, err := c.client.Pull(ctx, ref, pullOpts...)
 	if err != nil {
@@ -422,3 +426,42 @@ func (c *criService) encryptedImagesPullOpts() []containerd.RemoteOpt {
 	}
 	return nil
 }
+
+const (
+	// targetRefLabel is a label which contains image reference and will be passed
+	// to snapshotters.
+	targetRefLabel = "containerd.io/snapshot/cri.image-ref"
+	// targetDigestLabel is a label which contains layer digest and will be passed
+	// to snapshotters.
+	targetDigestLabel = "containerd.io/snapshot/cri.layer-digest"
+)
+
+// appendInfoHandlerWrapper makes a handler which appends some basic information
+// of images to each layer descriptor as annotations during unpack. These
+// annotations will be passed to snapshotters as labels. These labels will be
+// used mainly by stargz-based snapshotters for querying image contents from the
+// registry.
+func appendInfoHandlerWrapper(ref string) func(f containerdimages.Handler) containerdimages.Handler {
+	return func(f containerdimages.Handler) containerdimages.Handler {
+		return containerdimages.HandlerFunc(func(ctx context.Context, desc imagespec.Descriptor) ([]imagespec.Descriptor, error) {
+			children, err := f.Handle(ctx, desc)
+			if err != nil {
+				return nil, err
+			}
+			switch desc.MediaType {
+			case imagespec.MediaTypeImageManifest, containerdimages.MediaTypeDockerSchema2Manifest:
+				for i := range children {
+					c := &children[i]
+					if containerdimages.IsLayerType(c.MediaType) {
+						if c.Annotations == nil {
+							c.Annotations = make(map[string]string)
+						}
+						c.Annotations[targetRefLabel] = ref
+						c.Annotations[targetDigestLabel] = c.Digest.String()
+					}
+				}
+			}
+			return children, nil
+		})
+	}
+}