Add usage function to client

The usage function allows more configurable and accurate calculations
of the usage for an image.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>

image.go

@@ -19,6 +19,8 @@ package containerd
 import (
 	"context"
 	"fmt"
+	"strings"
+	"sync/atomic"
 
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/errdefs"
@@ -29,6 +31,7 @@ import (
 	"github.com/opencontainers/image-spec/identity"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
+	"golang.org/x/sync/semaphore"
 )
 
 // Image describes an image used by containers
@@ -45,6 +48,8 @@ type Image interface {
 	RootFS(ctx context.Context) ([]digest.Digest, error)
 	// Size returns the total size of the image's packed resources.
 	Size(ctx context.Context) (int64, error)
+	// Usage returns a usage calculation for the image.
+	Usage(context.Context, ...UsageOpt) (int64, error)
 	// Config descriptor for the image.
 	Config(ctx context.Context) (ocispec.Descriptor, error)
 	// IsUnpacked returns whether or not an image is unpacked.
@@ -53,6 +58,49 @@ type Image interface {
 	ContentStore() content.Store
 }
 
+type usageOptions struct {
+	manifestLimit *int
+	manifestOnly  bool
+	snapshots     bool
+}
+
+// UsageOpt is used to configure the usage calculation
+type UsageOpt func(*usageOptions) error
+
+// WithUsageManifestLimit sets the limit to the number of manifests which will
+// be walked for usage. Setting this value to 0 will require all manifests to
+// be walked, returning ErrNotFound if manifests are missing.
+// NOTE: By default all manifests which exist will be walked
+// and any non-existent manifests and their subobjects will be ignored.
+func WithUsageManifestLimit(i int) UsageOpt {
+	// If 0 then don't filter any manifests
+	// By default limits to current platform
+	return func(o *usageOptions) error {
+		o.manifestLimit = &i
+		return nil
+	}
+}
+
+// WithSnapshotUsage will check for referenced snapshots from the image objects
+// and include the snapshot size in the total usage.
+func WithSnapshotUsage() UsageOpt {
+	return func(o *usageOptions) error {
+		o.snapshots = true
+		return nil
+	}
+}
+
+// WithManifestUsage is used to get the usage for an image based on what is
+// reported by the manifests rather than what exists in the content store.
+// NOTE: This function is best used with the manifest limit set to get a
+// consistent value, otherwise non-existent manifests will be excluded.
+func WithManifestUsage() UsageOpt {
+	return func(o *usageOptions) error {
+		o.manifestOnly = true
+		return nil
+	}
+}
+
 var _ = (Image)(&image{})
 
 // NewImage returns a client image object from the metadata image
@@ -98,8 +146,95 @@ func (i *image) RootFS(ctx context.Context) ([]digest.Digest, error) {
 }
 
 func (i *image) Size(ctx context.Context) (int64, error) {
-	provider := i.client.ContentStore()
-	return i.i.Size(ctx, provider, i.platform)
+	return i.Usage(ctx, WithUsageManifestLimit(1), WithManifestUsage())
+}
+
+func (i *image) Usage(ctx context.Context, opts ...UsageOpt) (int64, error) {
+	var config usageOptions
+	for _, opt := range opts {
+		if err := opt(&config); err != nil {
+			return 0, err
+		}
+	}
+
+	var (
+		provider  = i.client.ContentStore()
+		handler   = images.ChildrenHandler(provider)
+		size      int64
+		mustExist bool
+	)
+
+	if config.manifestLimit != nil {
+		handler = images.LimitManifests(handler, i.platform, *config.manifestLimit)
+		mustExist = true
+	}
+
+	var wh images.HandlerFunc = func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		var usage int64
+		children, err := handler(ctx, desc)
+		if err != nil {
+			if !errdefs.IsNotFound(err) || mustExist {
+				return nil, err
+			}
+			if !config.manifestOnly {
+				// Do not count size of non-existent objects
+				desc.Size = 0
+			}
+		} else if config.snapshots || !config.manifestOnly {
+			info, err := provider.Info(ctx, desc.Digest)
+			if err != nil {
+				if !errdefs.IsNotFound(err) {
+					return nil, err
+				}
+				if !config.manifestOnly {
+					// Do not count size of non-existent objects
+					desc.Size = 0
+				}
+			} else if info.Size > desc.Size {
+				// Count actual usage, Size may be unset or -1
+				desc.Size = info.Size
+			}
+
+			for k, v := range info.Labels {
+				const prefix = "containerd.io/gc.ref.snapshot."
+				if !strings.HasPrefix(k, prefix) {
+					continue
+				}
+
+				sn := i.client.SnapshotService(k[len(prefix):])
+				if sn == nil {
+					continue
+				}
+
+				u, err := sn.Usage(ctx, v)
+				if err != nil {
+					if !errdefs.IsNotFound(err) && !errdefs.IsInvalidArgument(err) {
+						return nil, err
+					}
+				} else {
+					usage += u.Size
+				}
+			}
+		}
+
+		// Ignore unknown sizes. Generally unknown sizes should
+		// never be set in manifests, however, the usage
+		// calculation does not need to enforce this.
+		if desc.Size >= 0 {
+			usage += desc.Size
+		}
+
+		atomic.AddInt64(&size, usage)
+
+		return children, nil
+	}
+
+	l := semaphore.NewWeighted(3)
+	if err := images.Dispatch(ctx, wh, l, i.i.Target); err != nil {
+		return 0, err
+	}
+
+	return size, nil
 }
 
 func (i *image) Config(ctx context.Context) (ocispec.Descriptor, error) {

image_test.go

@@ -135,3 +135,101 @@ func TestImagePullWithDistSourceLabel(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestImageUsage(t *testing.T) {
+	if testing.Short() || runtime.GOOS == "windows" {
+		t.Skip()
+	}
+
+	imageName := "docker.io/library/busybox:latest"
+	ctx, cancel := testContext(t)
+	defer cancel()
+
+	client, err := newClient(t, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	// Cleanup
+	err = client.ImageService().Delete(ctx, imageName)
+	if err != nil && !errdefs.IsNotFound(err) {
+		t.Fatal(err)
+	}
+
+	testPlatform := platforms.Only(ocispec.Platform{
+		OS:           "linux",
+		Architecture: "amd64",
+	})
+
+	// Pull single platform, do not unpack
+	image, err := client.Pull(ctx, imageName, WithPlatformMatcher(testPlatform))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s1, err := image.Usage(ctx, WithUsageManifestLimit(1))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := image.Usage(ctx, WithUsageManifestLimit(0), WithManifestUsage()); err == nil {
+		t.Fatal("expected NotFound with missing manifests")
+	} else if !errdefs.IsNotFound(err) {
+		t.Fatalf("unexpected error: %+v", err)
+	}
+
+	// Pin image name to specific version for future fetches
+	imageName = imageName + "@" + image.Target().Digest.String()
+
+	// Fetch single platforms, but all manifests pulled
+	if _, err := client.Fetch(ctx, imageName, WithPlatformMatcher(testPlatform)); err != nil {
+		t.Fatal(err)
+	}
+
+	if s, err := image.Usage(ctx, WithUsageManifestLimit(1)); err != nil {
+		t.Fatal(err)
+	} else if s != s1 {
+		t.Fatalf("unexpected usage %d, expected %d", s, s1)
+	}
+
+	s2, err := image.Usage(ctx, WithUsageManifestLimit(0))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if s2 <= s1 {
+		t.Fatalf("Expected larger usage counting all manifests: %d <= %d", s2, s1)
+	}
+
+	s3, err := image.Usage(ctx, WithUsageManifestLimit(0), WithManifestUsage())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if s3 <= s2 {
+		t.Fatalf("Expected larger usage counting all manifest reported sizes: %d <= %d", s3, s2)
+	}
+
+	// Fetch everything
+	if _, err = client.Fetch(ctx, imageName); err != nil {
+		t.Fatal(err)
+	}
+
+	if s, err := image.Usage(ctx); err != nil {
+		t.Fatal(err)
+	} else if s != s3 {
+		t.Fatalf("Expected actual usage to equal manifest reported usage of %d: got %d", s3, s)
+	}
+
+	err = image.Unpack(ctx, DefaultSnapshotter)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if s, err := image.Usage(ctx, WithSnapshotUsage()); err != nil {
+		t.Fatal(err)
+	} else if s <= s3 {
+		t.Fatalf("Expected actual usage with snapshots to be greater: %d <= %d", s, s3)
+	}
+}