From 0c9b05fa60e9b3a8ab2b0eb0254833741a73db74 Mon Sep 17 00:00:00 2001
From: Phil Estes <estesp@linux.vnet.ibm.com>
Date: Wed, 20 May 2020 08:44:05 -0400
Subject: [PATCH] Fix image usage calculation error

Including snapshotter usage in total calculation should be gated by the
option `snapshotter` boolean.

Signed-off-by: Phil Estes <estesp@linux.vnet.ibm.com>
---
 image.go      | 36 +++++++++++++++++++-----------------
 image_test.go |  1 +
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/image.go b/image.go
index 174b008ea..c96b79301 100644
--- a/image.go
+++ b/image.go
@@ -203,24 +203,26 @@ func (i *image) Usage(ctx context.Context, opts ...UsageOpt) (int64, error) {
 				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
+			if config.snapshots {
+				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
 					}
-				} else {
-					usage += u.Size
 				}
 			}
 		}
diff --git a/image_test.go b/image_test.go
index 6d40b6ef1..38aa5fbb3 100644
--- a/image_test.go
+++ b/image_test.go
@@ -180,6 +180,7 @@ func TestImageUsage(t *testing.T) {
 
 	// Pin image name to specific version for future fetches
 	imageName = imageName + "@" + image.Target().Digest.String()
+	defer client.ImageService().Delete(ctx, imageName, images.SynchronousDelete())
 
 	// Fetch single platforms, but all manifests pulled
 	if _, err := client.Fetch(ctx, imageName, WithPlatformMatcher(testPlatform), WithAllMetadata()); err != nil {