cachemanager: size reporting in diskusage

Signed-off-by: Tonis Tiigi <[email protected]>

cache/manager.go

@@ -6,6 +6,8 @@ import (
 	"path/filepath"
 	"sync"
 
+	"golang.org/x/sync/errgroup"
+
 	"github.com/boltdb/bolt"
 	cdsnapshot "github.com/containerd/containerd/snapshot"
 	"github.com/pkg/errors"
@@ -33,7 +35,7 @@ type Accessor interface {
 }
 
 type Controller interface {
-	DiskUsage(ctx context.Context) ([]UsageInfo, error)
+	DiskUsage(ctx context.Context) ([]*UsageInfo, error)
 	Prune(ctx context.Context) (map[string]int64, error)
 	GC(ctx context.Context) error
 }
@@ -128,6 +130,7 @@ func (cm *cacheManager) get(id string) (ImmutableRef, error) {
 			cm:     cm,
 			refs:   make(map[*cacheRef]struct{}),
 			parent: parent,
+			size:   sizeUnknown,
 		}
 		cm.records[id] = rec // TODO: store to db
 	}
@@ -172,6 +175,7 @@ func (cm *cacheManager) New(s ImmutableRef) (MutableRef, error) {
 		cm:      cm,
 		refs:    make(map[*cacheRef]struct{}),
 		parent:  parent,
+		size:    sizeUnknown,
 	}
 
 	cm.mu.Lock()
@@ -203,23 +207,52 @@ func (cm *cacheManager) GetMutable(id string) (MutableRef, error) { // Rebase?
 	return rec.ref(), nil
 }
 
-func (cm *cacheManager) DiskUsage(ctx context.Context) ([]UsageInfo, error) {
+func (cm *cacheManager) DiskUsage(ctx context.Context) ([]*UsageInfo, error) {
 	cm.mu.Lock()
-	defer cm.mu.Unlock()
 
-	var du []UsageInfo
+	var du []*UsageInfo
 
 	for id, cr := range cm.records {
 		cr.mu.Lock()
-		c := UsageInfo{
+		c := &UsageInfo{
 			ID:     id,
 			Active: cr.mutable,
 			InUse:  len(cr.refs) > 0,
-			Size:   -1, // TODO:
+			Size:   cr.size,
+		}
+		if cr.mutable && len(cr.refs) > 0 && !cr.frozen {
+			c.Size = 0 // size can not be determined because it is changing
 		}
 		cr.mu.Unlock()
 		du = append(du, c)
 	}
+	cm.mu.Unlock()
+
+	eg, ctx := errgroup.WithContext(ctx)
+
+	for _, d := range du {
+		if d.Size == sizeUnknown {
+			func(d *UsageInfo) {
+				eg.Go(func() error {
+					ref, err := cm.Get(d.ID)
+					if err != nil {
+						d.Size = 0
+						return nil
+					}
+					s, err := ref.Size(ctx)
+					if err != nil {
+						return err
+					}
+					d.Size = s
+					return ref.Release()
+				})
+			}(d)
+		}
+	}
+
+	if err := eg.Wait(); err != nil {
+		return du, err
+	}
 
 	return du, nil
 }

cache/refs.go

@@ -1,20 +1,23 @@
 package cache
 
 import (
-	"context"
 	"crypto/rand"
 	"encoding/hex"
 	"sync"
 
 	"github.com/containerd/containerd/mount"
 	"github.com/pkg/errors"
+	"github.com/tonistiigi/buildkit_poc/util/flightcontrol"
+	"golang.org/x/net/context"
 )
 
+const sizeUnknown int64 = -1
+
 type ImmutableRef interface {
 	Mountable
 	ID() string
 	Release() error
-	Size() (int64, error)
+	Size(ctx context.Context) (int64, error)
 	// Prepare() / ChainID() / Meta()
 }
 
@@ -23,7 +26,7 @@ type MutableRef interface {
 	ID() string
 	Freeze() (ImmutableRef, error)
 	ReleaseAndCommit(ctx context.Context) (ImmutableRef, error)
-	Size() (int64, error)
+	Size(ctx context.Context) (int64, error)
 }
 
 type Mountable interface {
@@ -41,6 +44,9 @@ type cacheRecord struct {
 	parent    ImmutableRef
 	view      string
 	viewMount []mount.Mount
+
+	sizeG flightcontrol.Group
+	size  int64
 }
 
 // hold manager lock before calling
@@ -50,6 +56,27 @@ func (cr *cacheRecord) ref() *cacheRef {
 	return ref
 }
 
+func (cr *cacheRecord) Size(ctx context.Context) (int64, error) {
+	// this expects that usage() is implemented lazily
+	s, err, _ := cr.sizeG.Do(ctx, cr.id, func(ctx context.Context) (interface{}, error) {
+		cr.mu.Lock()
+		s := cr.size
+		cr.mu.Unlock()
+		if s != sizeUnknown {
+			return s, nil
+		}
+		usage, err := cr.cm.ManagerOpt.Snapshotter.Usage(ctx, cr.id)
+		if err != nil {
+			return s, errors.Wrapf(err, "failed to get usage for %s", cr.id)
+		}
+		cr.mu.Lock()
+		cr.size = s
+		cr.mu.Unlock()
+		return usage.Size, nil
+	})
+	return s.(int64), err
+}
+
 type cacheRef struct {
 	*cacheRecord
 }
@@ -130,6 +157,7 @@ func (sr *cacheRef) Freeze() (ImmutableRef, error) {
 	}
 
 	sr.frozen = true
+	sr.size = sizeUnknown
 
 	return sr, nil
 }
@@ -164,16 +192,13 @@ func (sr *cacheRef) ReleaseAndCommit(ctx context.Context) (ImmutableRef, error)
 		id:   id,
 		cm:   sr.cm,
 		refs: make(map[*cacheRef]struct{}),
+		size: sizeUnknown,
 	}
 	sr.cm.records[id] = rec // TODO: save to db
 
 	return rec.ref(), nil
 }
 
-func (sr *cacheRef) Size() (int64, error) {
-	return -1, errors.New("Size not implemented")
-}
-
 func (sr *cacheRef) ID() string {
 	return sr.id
 }

control/control_test.go

@@ -16,7 +16,7 @@ import (
 	"github.com/containerd/containerd/mount"
 	"github.com/containerd/containerd/rootfs"
 	cdsnapshot "github.com/containerd/containerd/snapshot"
-	"github.com/containerd/containerd/snapshot/naive"
+	"github.com/containerd/containerd/snapshot/overlay"
 	digest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
@@ -89,8 +89,23 @@ func TestControl(t *testing.T) {
 	lm.Unmount()
 	assert.NoError(t, err)
 
+	du, err := cm.DiskUsage(context.TODO())
+	assert.NoError(t, err)
+
+	// fmt.Printf("du1:\n")
+	// for _, d := range du {
+	// 	fmt.Printf("du1: %+v\n", d)
+	// }
+
 	err = snap.Release()
 	assert.NoError(t, err)
+
+	du, err = cm.DiskUsage(context.TODO())
+	assert.NoError(t, err)
+
+	for _, d := range du {
+		assert.True(t, d.Size >= 8192)
+	}
 }
 
 type containerd struct {
@@ -100,7 +115,7 @@ type containerd struct {
 }
 
 func localContainerd(root string) (*containerd, error) {
-	s, err := naive.NewSnapshotter(filepath.Join(root, "snapshots"))
+	s, err := overlay.NewSnapshotter(filepath.Join(root, "snapshots"))
 	if err != nil {
 		return nil, err
 	}
@@ -131,7 +146,6 @@ func (a *localApplier) Apply(ctx context.Context, desc ocispec.Descriptor, mount
 		return ocispec.Descriptor{}, errors.Wrap(err, "failed to create temporary directory")
 	}
 	defer os.RemoveAll(dir)
-
 	if err := mount.MountAll(mounts, dir); err != nil {
 		return ocispec.Descriptor{}, errors.Wrap(err, "failed to mount")
 	}

hack/test

@@ -4,4 +4,4 @@ set -eu -o pipefail -x
 
 # update this to iidfile after 17.06
 docker build -t buildkit_poc:test -f ./hack/dockerfiles/test.Dockerfile --force-rm .
-docker run --cap-add=SYS_ADMIN buildkit_poc:test go test ${TESTFLAGS:--v} ${TESTPKGS:-./...}
+docker run -v /tmp --cap-add=SYS_ADMIN buildkit_poc:test go test ${TESTFLAGS:--v} ${TESTPKGS:-./...}

snapshot/localmounter.go

@@ -44,7 +44,7 @@ func (lm *localMounter) Mount() (string, error) {
 		return "", err
 	}
 	lm.target = dir
-	return "", nil
+	return dir, nil
 }
 
 func (lm *localMounter) Unmount() error {