From 2432b54a56050498f3e8ceef5ebeddf5ab88485f Mon Sep 17 00:00:00 2001
From: James Jenkins <James.Jenkins@ibm.com>
Date: Wed, 21 Sep 2022 13:21:28 -0400
Subject: [PATCH] Add new ctr option for discarding unpacked layers

Add a new ctr cli option, allowing the garbage collector to discard any
unpacked layers after importing an image. This new option is
incompatible with the no-unpack ctr import option.

Signed-off-by: James Jenkins <James.Jenkins@ibm.com>
---
 cmd/ctr/commands/images/import.go | 18 +++++++++++++++++-
 import.go                         | 16 +++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/cmd/ctr/commands/images/import.go b/cmd/ctr/commands/images/import.go
index 850b8a27c..fa1f1d801 100644
--- a/cmd/ctr/commands/images/import.go
+++ b/cmd/ctr/commands/images/import.go
@@ -79,12 +79,16 @@ If foobar.tar contains an OCI ref named "latest" and anonymous ref "sha256:deadb
 		},
 		cli.BoolFlag{
 			Name:  "no-unpack",
-			Usage: "skip unpacking the images, false by default",
+			Usage: "skip unpacking the images, cannot be used with --discard-unpacked-layers, false by default",
 		},
 		cli.BoolFlag{
 			Name:  "compress-blobs",
 			Usage: "compress uncompressed blobs when creating manifest (Docker format only)",
 		},
+		cli.BoolFlag{
+			Name:  "discard-unpacked-layers",
+			Usage: "allow the garbage collector to clean layers up from the content store after unpacking, cannot be used with --no-unpack, false by default",
+		},
 	}, commands.SnapshotterFlags...),
 
 	Action: func(context *cli.Context) error {
@@ -132,12 +136,23 @@ If foobar.tar contains an OCI ref named "latest" and anonymous ref "sha256:deadb
 
 		opts = append(opts, containerd.WithAllPlatforms(context.Bool("all-platforms")))
 
+		if context.Bool("discard-unpacked-layers") && context.Bool("no-unpack") {
+			return fmt.Errorf("--discard-unpacked-layers and --no-unpack are incompatible options")
+		}
+		opts = append(opts, containerd.WithDiscardUnpackedLayers(context.Bool("discard-unpacked-layers")))
+
 		client, ctx, cancel, err := commands.NewClient(context)
 		if err != nil {
 			return err
 		}
 		defer cancel()
 
+		ctx, done, err := client.WithLease(ctx)
+		if err != nil {
+			return err
+		}
+		defer done(ctx)
+
 		var r io.ReadCloser
 		if in == "-" {
 			r = os.Stdin
@@ -147,6 +162,7 @@ If foobar.tar contains an OCI ref named "latest" and anonymous ref "sha256:deadb
 				return err
 			}
 		}
+
 		imgs, err := client.Import(ctx, r, opts...)
 		closeErr := r.Close()
 		if err != nil {
diff --git a/import.go b/import.go
index 8936d88ef..98e94df1f 100644
--- a/import.go
+++ b/import.go
@@ -38,6 +38,7 @@ type importOpts struct {
 	allPlatforms    bool
 	platformMatcher platforms.MatchComparer
 	compress        bool
+	discardLayers   bool
 }
 
 // ImportOpt allows the caller to specify import specific options
@@ -105,6 +106,15 @@ func WithImportCompression() ImportOpt {
 	}
 }
 
+// WithDiscardUnpackedLayers allows the garbage collector to clean up
+// layers from content store after unpacking.
+func WithDiscardUnpackedLayers(discard bool) ImportOpt {
+	return func(c *importOpts) error {
+		c.discardLayers = discard
+		return nil
+	}
+}
+
 // Import imports an image from a Tar stream using reader.
 // Caller needs to specify importer. Future version may use oci.v1 as the default.
 // Note that unreferenced blobs may be imported to the content store as well.
@@ -195,7 +205,11 @@ func (c *Client) Import(ctx context.Context, reader io.Reader, opts ...ImportOpt
 	}
 
 	handler = images.FilterPlatforms(handler, platformMatcher)
-	handler = images.SetChildrenLabels(cs, handler)
+	if iopts.discardLayers {
+		handler = images.SetChildrenMappedLabels(cs, handler, images.ChildGCLabelsFilterLayers)
+	} else {
+		handler = images.SetChildrenLabels(cs, handler)
+	}
 	if err := images.WalkNotEmpty(ctx, handler, index); err != nil {
 		return nil, err
 	}