From f4695e7b4e1ec802d1196ff9cf3398783c5920d3 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Wed, 31 May 2017 16:52:53 -0700
Subject: [PATCH] blobmapping: account blob size in disk usage

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 control/control_test.go             |  3 +--
 snapshot/blobmapping/snapshotter.go | 21 ++++++++++++++++++++-
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/control/control_test.go b/control/control_test.go
index eb9759ebcf0b..2e22e41649ea 100644
--- a/control/control_test.go
+++ b/control/control_test.go
@@ -92,9 +92,8 @@ func TestControl(t *testing.T) {
 	du, err := cm.DiskUsage(context.TODO())
 	assert.NoError(t, err)
 
-	// fmt.Printf("du1:\n")
 	// for _, d := range du {
-	// 	fmt.Printf("du1: %+v\n", d)
+	// 	fmt.Printf("du: %+v\n", d)
 	// }
 
 	err = snap.Release()
diff --git a/snapshot/blobmapping/snapshotter.go b/snapshot/blobmapping/snapshotter.go
index b6bc377bc12f..b6ab7499c9b2 100644
--- a/snapshot/blobmapping/snapshotter.go
+++ b/snapshot/blobmapping/snapshotter.go
@@ -88,13 +88,32 @@ func (s *Snapshotter) Remove(ctx context.Context, key string) error {
 			b = tx.Bucket(bucketByBlob)
 			b.Delete(blobKey(blob, key))
 			if len(keyRange(tx, blobKey(blob, ""))) == 0 { // last snapshot
-				s.opt.Content.Delete(ctx, digest.Digest(blob)) // log error
+				s.opt.Content.Delete(ctx, blob) // log error
 			}
 		}
 		return nil
 	})
 }
 
+func (s *Snapshotter) Usage(ctx context.Context, key string) (snapshot.Usage, error) {
+	u, err := s.Snapshotter.Usage(ctx, key)
+	if err != nil {
+		return snapshot.Usage{}, err
+	}
+	blob, err := s.GetBlob(ctx, key)
+	if err != nil {
+		return u, err
+	}
+	if blob != "" {
+		info, err := s.opt.Content.Info(ctx, blob)
+		if err != nil {
+			return u, err
+		}
+		(&u).Add(snapshot.Usage{Size: info.Size, Inodes: 1})
+	}
+	return u, nil
+}
+
 // TODO: make Blob/SetBlob part of generic metadata wrapper that can detect
 // blob key for deletion logic