Set content labels based on content type

dmcgowan · dmcgowan · commit c8b14ae4c01e · 2020-07-21T00:46:07.000-07:00
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 &lt;derek@mcg.dev&gt;
diff --git a/client.go b/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 {
diff --git a/client_opts.go b/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.
diff --git a/client_test.go b/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 {
diff --git a/images/handlers.go b/images/handlers.go
@@ -170,21 +170,46 @@ 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 {
 			return children, err
 		}
 
 		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...)
diff --git a/images/mediatypes.go b/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)
+}
diff --git a/pull.go b/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
diff --git a/unpacker.go b/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,20 +312,18 @@ func (u *unpacker) handlerWrapper(
 				lock.Lock()
 				for _, nl := range nonLayers {
 					layers[nl.Digest] = manifestLayers
-					parent[nl.Digest] = desc
 				}
 				lock.Unlock()
 
 				children = nonLayers
 			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)
 					})
 				}
 			}
