Support cancellation in `directory.Size()`

Makes sure that if the user cancels a request that the daemon stops
trying to traverse a directory.

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

api/server/router/system/system_routes.go

@@ -73,7 +73,7 @@ func (s *systemRouter) getDiskUsage(ctx context.Context, w http.ResponseWriter,
 	if err != nil {
 		return err
 	}
-	builderSize, err := s.builder.DiskUsage()
+	builderSize, err := s.builder.DiskUsage(ctx)
 	if err != nil {
 		return pkgerrors.Wrap(err, "error getting build cache usage")
 	}

builder/fscache/fscache.go

@@ -154,8 +154,8 @@ func (fsc *FSCache) SyncFrom(ctx context.Context, id RemoteIdentifier) (builder.
 }
 
 // DiskUsage reports how much data is allocated by the cache
-func (fsc *FSCache) DiskUsage() (int64, error) {
-	return fsc.store.DiskUsage()
+func (fsc *FSCache) DiskUsage(ctx context.Context) (int64, error) {
+	return fsc.store.DiskUsage(ctx)
 }
 
 // Prune allows manually cleaning up the cache
@@ -382,14 +382,14 @@ func (s *fsCacheStore) Get(id string) (*cachedSourceRef, error) {
 }
 
 // DiskUsage reports how much data is allocated by the cache
-func (s *fsCacheStore) DiskUsage() (int64, error) {
+func (s *fsCacheStore) DiskUsage(ctx context.Context) (int64, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	var size int64
 
 	for _, snap := range s.sources {
 		if len(snap.refs) == 0 {
-			ss, err := snap.getSize()
+			ss, err := snap.getSize(ctx)
 			if err != nil {
 				return 0, err
 			}
@@ -414,7 +414,7 @@ func (s *fsCacheStore) Prune(ctx context.Context) (uint64, error) {
 		default:
 		}
 		if len(snap.refs) == 0 {
-			ss, err := snap.getSize()
+			ss, err := snap.getSize(ctx)
 			if err != nil {
 				return size, err
 			}
@@ -433,6 +433,7 @@ func (s *fsCacheStore) GC() error {
 	defer s.mu.Unlock()
 	var size uint64
 
+	ctx := context.Background()
 	cutoff := time.Now().Add(-s.gcPolicy.MaxKeepDuration)
 	var blacklist []*cachedSource
 
@@ -443,7 +444,7 @@ func (s *fsCacheStore) GC() error {
 					return errors.Wrapf(err, "failed to delete %s", id)
 				}
 			} else {
-				ss, err := snap.getSize()
+				ss, err := snap.getSize(ctx)
 				if err != nil {
 					return err
 				}
@@ -458,7 +459,7 @@ func (s *fsCacheStore) GC() error {
 		if size <= s.gcPolicy.MaxSize {
 			break
 		}
-		ss, err := snap.getSize()
+		ss, err := snap.getSize(ctx)
 		if err != nil {
 			return err
 		}
@@ -521,9 +522,9 @@ func (cs *cachedSource) getRef() *cachedSourceRef {
 }
 
 // hold storage lock before calling
-func (cs *cachedSource) getSize() (int64, error) {
+func (cs *cachedSource) getSize(ctx context.Context) (int64, error) {
 	if cs.sourceMeta.Size < 0 {
-		ss, err := directory.Size(cs.dir)
+		ss, err := directory.Size(ctx, cs.dir)
 		if err != nil {
 			return 0, err
 		}

builder/fscache/fscache_test.go

@@ -59,13 +59,13 @@ func TestFSCache(t *testing.T) {
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(string(dt), "data2"))
 
-	s, err := fscache.DiskUsage()
+	s, err := fscache.DiskUsage(context.TODO())
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(s, int64(0)))
 
 	assert.Check(t, src3.Close())
 
-	s, err = fscache.DiskUsage()
+	s, err = fscache.DiskUsage(context.TODO())
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(s, int64(5)))
 
@@ -80,7 +80,7 @@ func TestFSCache(t *testing.T) {
 	assert.Check(t, is.Equal(src4.Root().Path(), src3.Root().Path()))
 	assert.Check(t, src4.Close())
 
-	s, err = fscache.DiskUsage()
+	s, err = fscache.DiskUsage(context.TODO())
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(s, int64(10)))
 
@@ -93,7 +93,7 @@ func TestFSCache(t *testing.T) {
 	time.Sleep(100 * time.Millisecond)
 
 	// only last insertion after GC
-	s, err = fscache.DiskUsage()
+	s, err = fscache.DiskUsage(context.TODO())
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(s, int64(8)))
 
@@ -102,7 +102,7 @@ func TestFSCache(t *testing.T) {
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(released, uint64(8)))
 
-	s, err = fscache.DiskUsage()
+	s, err = fscache.DiskUsage(context.TODO())
 	assert.Check(t, err)
 	assert.Check(t, is.Equal(s, int64(0)))
 }

daemon/disk_usage.go

@@ -53,7 +53,7 @@ func (daemon *Daemon) SystemDiskUsage(ctx context.Context) (*types.DiskUsage, er
 			refs := daemon.volumes.Refs(v)
 
 			tv := volumeToAPIType(v)
-			sz, err := directory.Size(v.Path())
+			sz, err := directory.Size(ctx, v.Path())
 			if err != nil {
 				logrus.Warnf("failed to determine size of volume %v", name)
 				sz = -1

daemon/graphdriver/aufs/aufs.go

@@ -24,6 +24,7 @@ package aufs // import "github.com/docker/docker/daemon/graphdriver/aufs"
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -502,7 +503,7 @@ func (a *Driver) DiffSize(id, parent string) (size int64, err error) {
 		return a.naiveDiff.DiffSize(id, parent)
 	}
 	// AUFS doesn't need the parent layer to calculate the diff size.
-	return directory.Size(path.Join(a.rootPath(), "diff", id))
+	return directory.Size(context.TODO(), path.Join(a.rootPath(), "diff", id))
 }
 
 // ApplyDiff extracts the changeset from the given diff into the

daemon/graphdriver/overlay2/overlay.go

@@ -4,6 +4,7 @@ package overlay2 // import "github.com/docker/docker/daemon/graphdriver/overlay2
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -706,7 +707,7 @@ func (d *Driver) ApplyDiff(id string, parent string, diff io.Reader) (size int64
 		return 0, err
 	}
 
-	return directory.Size(applyDir)
+	return directory.Size(context.TODO(), applyDir)
 }
 
 func (d *Driver) getDiffPath(id string) string {
@@ -722,7 +723,7 @@ func (d *Driver) DiffSize(id, parent string) (size int64, err error) {
 	if useNaiveDiff(d.home) || !d.isParent(id, parent) {
 		return d.naiveDiff.DiffSize(id, parent)
 	}
-	return directory.Size(d.getDiffPath(id))
+	return directory.Size(context.TODO(), d.getDiffPath(id))
 }
 
 // Diff produces an archive of the changes between the specified

daemon/prune.go

@@ -125,7 +125,7 @@ func (daemon *Daemon) VolumesPrune(ctx context.Context, pruneFilters filters.Arg
 					return nil
 				}
 			}
-			vSize, err := directory.Size(v.Path())
+			vSize, err := directory.Size(ctx, v.Path())
 			if err != nil {
 				logrus.Warnf("could not determine size of volume %s: %v", name, err)
 			}

pkg/directory/directory_test.go

@@ -1,6 +1,7 @@
 package directory // import "github.com/docker/docker/pkg/directory"
 
 import (
+	"context"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -18,7 +19,7 @@ func TestSizeEmpty(t *testing.T) {
 	}
 
 	var size int64
-	if size, _ = Size(dir); size != 0 {
+	if size, _ = Size(context.Background(), dir); size != 0 {
 		t.Fatalf("empty directory has size: %d", size)
 	}
 }
@@ -37,7 +38,7 @@ func TestSizeEmptyFile(t *testing.T) {
 	}
 
 	var size int64
-	if size, _ = Size(file.Name()); size != 0 {
+	if size, _ = Size(context.Background(), file.Name()); size != 0 {
 		t.Fatalf("directory with one file has size: %d", size)
 	}
 }
@@ -59,7 +60,7 @@ func TestSizeNonemptyFile(t *testing.T) {
 	file.Write(d)
 
 	var size int64
-	if size, _ = Size(file.Name()); size != 5 {
+	if size, _ = Size(context.Background(), file.Name()); size != 5 {
 		t.Fatalf("directory with one 5-byte file has size: %d", size)
 	}
 }
@@ -76,7 +77,7 @@ func TestSizeNestedDirectoryEmpty(t *testing.T) {
 	}
 
 	var size int64
-	if size, _ = Size(dir); size != 0 {
+	if size, _ = Size(context.Background(), dir); size != 0 {
 		t.Fatalf("directory with one empty directory has size: %d", size)
 	}
 }
@@ -101,7 +102,7 @@ func TestSizeFileAndNestedDirectoryEmpty(t *testing.T) {
 	file.Write(d)
 
 	var size int64
-	if size, _ = Size(dir); size != 6 {
+	if size, _ = Size(context.Background(), dir); size != 6 {
 		t.Fatalf("directory with 6-byte file and empty directory has size: %d", size)
 	}
 }
@@ -134,7 +135,7 @@ func TestSizeFileAndNestedDirectoryNonempty(t *testing.T) {
 	nestedFile.Write(nestedData)
 
 	var size int64
-	if size, _ = Size(dir); size != 12 {
+	if size, _ = Size(context.Background(), dir); size != 12 {
 		t.Fatalf("directory with 6-byte file and nested directory with 6-byte file has size: %d", size)
 	}
 }
@@ -186,7 +187,7 @@ func TestMoveToSubdir(t *testing.T) {
 
 // Test a non-existing directory
 func TestSizeNonExistingDirectory(t *testing.T) {
-	if _, err := Size("/thisdirectoryshouldnotexist/TestSizeNonExistingDirectory"); err == nil {
+	if _, err := Size(context.Background(), "/thisdirectoryshouldnotexist/TestSizeNonExistingDirectory"); err == nil {
 		t.Fatalf("error is expected")
 	}
 }

pkg/directory/directory_unix.go

@@ -3,13 +3,14 @@
 package directory // import "github.com/docker/docker/pkg/directory"
 
 import (
+	"context"
 	"os"
 	"path/filepath"
 	"syscall"
 )
 
 // Size walks a directory tree and returns its total size in bytes.
-func Size(dir string) (size int64, err error) {
+func Size(ctx context.Context, dir string) (size int64, err error) {
 	data := make(map[uint64]struct{})
 	err = filepath.Walk(dir, func(d string, fileInfo os.FileInfo, err error) error {
 		if err != nil {
@@ -20,6 +21,11 @@ func Size(dir string) (size int64, err error) {
 			}
 			return err
 		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
 
 		// Ignore directory sizes
 		if fileInfo == nil {

pkg/directory/directory_windows.go

@@ -1,12 +1,13 @@
 package directory // import "github.com/docker/docker/pkg/directory"
 
 import (
+	"context"
 	"os"
 	"path/filepath"
 )
 
 // Size walks a directory tree and returns its total size in bytes.
-func Size(dir string) (size int64, err error) {
+func Size(ctx context.Context, dir string) (size int64, err error) {
 	err = filepath.Walk(dir, func(d string, fileInfo os.FileInfo, err error) error {
 		if err != nil {
 			// if dir does not exist, Size() returns the error.
@@ -17,6 +18,12 @@ func Size(dir string) (size int64, err error) {
 			return err
 		}
 
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		// Ignore directory sizes
 		if fileInfo == nil {
 			return nil