Set content labels based on content type

Give control of the content labeling process for children to
the client. This allows the client to control the names
associated with the labels and filter out labels.

Signed-off-by: Derek McGowan <derek@mcg.dev>

client.go

@@ -312,11 +312,6 @@ type RemoteContext struct {
 	// afterwards. Unpacking is required to run an image.
 	Unpack bool
 
-	// DiscardContent is a boolean flag to specify whether to allow GC to clean
-	// layers up from the content store after successfully unpacking these
-	// contents to the snapshotter.
-	DiscardContent bool
-
 	// UnpackOpts handles options to the unpack call.
 	UnpackOpts []UnpackOpt
 
@@ -356,6 +351,10 @@ type RemoteContext struct {
 
 	// AllMetadata downloads all manifests and known-configuration files
 	AllMetadata bool
+
+	// ChildLabelMap sets the labels used to reference child objects in the content
+	// store. By default, all GC reference labels will be set for all fetched content.
+	ChildLabelMap func(ocispec.Descriptor) []string
 }
 
 func defaultRemoteContext() *RemoteContext {

client_opts.go

@@ -23,6 +23,7 @@ import (
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/snapshots"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 
 	"google.golang.org/grpc"
 )
@@ -132,14 +133,6 @@ func WithPullUnpack(_ *Client, c *RemoteContext) error {
 	return nil
 }
 
-// WithDiscardContent is used to allow GC to clean layers up from
-// the content store after successfully unpacking these contents to
-// the snapshotter.
-func WithDiscardContent(_ *Client, c *RemoteContext) error {
-	c.DiscardContent = true
-	return nil
-}
-
 // WithUnpackOpts is used to add unpack options to the unpacker.
 func WithUnpackOpts(opts []UnpackOpt) RemoteOpt {
 	return func(_ *Client, c *RemoteContext) error {
@@ -183,6 +176,18 @@ func WithPullLabels(labels map[string]string) RemoteOpt {
 	}
 }
 
+// WithChildLabelMap sets the map function used to define the labels set
+// on referenced child content in the content store. This can be used
+// to overwrite the default GC labels or filter which labels get set
+// for content.
+// The default is `images.ChildGCLabels`.
+func WithChildLabelMap(fn func(ocispec.Descriptor) []string) RemoteOpt {
+	return func(_ *Client, c *RemoteContext) error {
+		c.ChildLabelMap = fn
+		return nil
+	}
+}
+
 // WithSchema1Conversion is used to convert Docker registry schema 1
 // manifests to oci manifests on pull. Without this option schema 1
 // manifests will return a not supported error.

client_test.go

@@ -229,6 +229,11 @@ func TestImagePullWithDiscardContent(t *testing.T) {
 	ctx, cancel := testContext(t)
 	defer cancel()
 
+	err = client.ImageService().Delete(ctx, testImage, images.SynchronousDelete())
+	if err != nil {
+		t.Fatal(err)
+	}
+
 	ls := client.LeasesService()
 	l, err := ls.Create(ctx, leases.WithRandomID(), leases.WithExpiration(24*time.Hour))
 	if err != nil {
@@ -238,7 +243,7 @@ func TestImagePullWithDiscardContent(t *testing.T) {
 	img, err := client.Pull(ctx, testImage,
 		WithPlatformMatcher(platforms.Default()),
 		WithPullUnpack,
-		WithDiscardContent,
+		WithChildLabelMap(images.ChildGCLabelsFilterLayers),
 	)
 	// Synchronously garbage collect contents
 	if errL := ls.Delete(ctx, l, leases.SynchronousDelete); errL != nil {

images/handlers.go

@@ -170,6 +170,19 @@ func ChildrenHandler(provider content.Provider) HandlerFunc {
 // the children returned by the handler and passes through the children.
 // Must follow a handler that returns the children to be labeled.
 func SetChildrenLabels(manager content.Manager, f HandlerFunc) HandlerFunc {
+	return SetChildrenMappedLabels(manager, f, nil)
+}
+
+// SetChildrenMappedLabels is a handler wrapper which sets labels for the content on
+// the children returned by the handler and passes through the children.
+// Must follow a handler that returns the children to be labeled.
+// The label map allows the caller to control the labels per child descriptor.
+// For returned labels, the index of the child will be appended to the end
+// except for the first index when the returned label does not end with '.'.
+func SetChildrenMappedLabels(manager content.Manager, f HandlerFunc, labelMap func(ocispec.Descriptor) []string) HandlerFunc {
+	if labelMap == nil {
+		labelMap = ChildGCLabels
+	}
 	return func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
 		children, err := f(ctx, desc)
 		if err != nil {
@@ -177,14 +190,26 @@ func SetChildrenLabels(manager content.Manager, f HandlerFunc) HandlerFunc {
 		}
 
 		if len(children) > 0 {
-			info := content.Info{
-				Digest: desc.Digest,
-				Labels: map[string]string{},
-			}
-			fields := []string{}
-			for i, ch := range children {
-				info.Labels[fmt.Sprintf("containerd.io/gc.ref.content.%d", i)] = ch.Digest.String()
-				fields = append(fields, fmt.Sprintf("labels.containerd.io/gc.ref.content.%d", i))
+			var (
+				info = content.Info{
+					Digest: desc.Digest,
+					Labels: map[string]string{},
+				}
+				fields = []string{}
+				keys   = map[string]uint{}
+			)
+			for _, ch := range children {
+				labelKeys := labelMap(ch)
+				for _, key := range labelKeys {
+					idx := keys[key]
+					keys[key] = idx + 1
+					if idx > 0 || key[len(key)-1] == '.' {
+						key = fmt.Sprintf("%s%d", key, idx)
+					}
+
+					info.Labels[key] = ch.Digest.String()
+					fields = append(fields, "labels."+key)
+				}
 			}
 
 			_, err := manager.Update(ctx, info, fields...)

images/mediatypes.go

@@ -124,3 +124,31 @@ func IsKnownConfig(mt string) bool {
 	}
 	return false
 }
+
+// ChildGCLabels returns the label for a given descriptor to reference it
+func ChildGCLabels(desc ocispec.Descriptor) []string {
+	mt := desc.MediaType
+	if IsKnownConfig(mt) {
+		return []string{"containerd.io/gc.ref.content.config"}
+	}
+
+	switch mt {
+	case MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
+		return []string{"containerd.io/gc.ref.content.m."}
+	}
+
+	if IsLayerType(mt) {
+		return []string{"containerd.io/gc.ref.content.l."}
+	}
+
+	return []string{"containerd.io/gc.ref.content."}
+}
+
+// ChildGCLabelsFilterLayers returns the labels for a given descriptor to
+// reference it, skipping layer media types
+func ChildGCLabelsFilterLayers(desc ocispec.Descriptor) []string {
+	if IsLayerType(desc.MediaType) {
+		return nil
+	}
+	return ChildGCLabels(desc)
+}

pull.go

@@ -159,7 +159,7 @@ func (c *Client) fetch(ctx context.Context, rCtx *RemoteContext, ref string, lim
 		// Get all the children for a descriptor
 		childrenHandler := images.ChildrenHandler(store)
 		// Set any children labels for that content
-		childrenHandler = images.SetChildrenLabels(store, childrenHandler)
+		childrenHandler = images.SetChildrenMappedLabels(store, childrenHandler, rCtx.ChildLabelMap)
 		if rCtx.AllMetadata {
 			// Filter manifests by platforms but allow to handle manifest
 			// and configuration for not-target platforms

unpacker.go

@@ -22,7 +22,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"math/rand"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -78,7 +77,6 @@ func (u *unpacker) unpack(
 	rCtx *RemoteContext,
 	h images.Handler,
 	config ocispec.Descriptor,
-	parentDesc ocispec.Descriptor,
 	layers []ocispec.Descriptor,
 ) error {
 	p, err := content.ReadBlob(ctx, u.c.ContentStore(), config)
@@ -247,31 +245,6 @@ EachLayer:
 		"chainID": chainID,
 	}).Debug("image unpacked")
 
-	if rCtx.DiscardContent {
-		// delete references to successfully unpacked layers
-		layersMap := map[string]struct{}{}
-		for _, desc := range layers {
-			layersMap[desc.Digest.String()] = struct{}{}
-		}
-		pinfo, err := cs.Info(ctx, parentDesc.Digest)
-		if err != nil {
-			return err
-		}
-		fields := []string{}
-		for k, v := range pinfo.Labels {
-			if strings.HasPrefix(k, "containerd.io/gc.ref.content.") {
-				if _, ok := layersMap[v]; ok {
-					// We've already unpacked this layer content
-					pinfo.Labels[k] = ""
-					fields = append(fields, "labels."+k)
-				}
-			}
-		}
-		if _, err := cs.Update(ctx, pinfo, fields...); err != nil {
-			return err
-		}
-	}
-
 	return nil
 }
 
@@ -314,7 +287,6 @@ func (u *unpacker) handlerWrapper(
 		var (
 			lock   sync.Mutex
 			layers = map[digest.Digest][]ocispec.Descriptor{}
-			parent = map[digest.Digest]ocispec.Descriptor{}
 		)
 		return images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
 			children, err := f.Handle(ctx, desc)
@@ -340,7 +312,6 @@ func (u *unpacker) handlerWrapper(
 				lock.Lock()
 				for _, nl := range nonLayers {
 					layers[nl.Digest] = manifestLayers
-					parent[nl.Digest] = desc
 				}
 				lock.Unlock()
 
@@ -348,12 +319,11 @@ func (u *unpacker) handlerWrapper(
 			case images.MediaTypeDockerSchema2Config, ocispec.MediaTypeImageConfig:
 				lock.Lock()
 				l := layers[desc.Digest]
-				p := parent[desc.Digest]
 				lock.Unlock()
 				if len(l) > 0 {
 					atomic.AddInt32(unpacks, 1)
 					eg.Go(func() error {
-						return u.unpack(uctx, rCtx, f, desc, p, l)
+						return u.unpack(uctx, rCtx, f, desc, l)
 					})
 				}
 			}