From cf7d246ab0d848cbc673f7c2d57031aecacacdfe Mon Sep 17 00:00:00 2001
From: Kenfe-Mickael Laventure <mickael.laventure@gmail.com>
Date: Wed, 12 Apr 2017 13:59:31 -0700
Subject: [PATCH] Honor context within SystemDiskUsage

Signed-off-by: Kenfe-Mickael Laventure <mickael.laventure@gmail.com>
---
 api/server/router/system/backend.go       |  2 +-
 api/server/router/system/system.go        |  2 +-
 api/server/router/system/system_routes.go |  2 +-
 daemon/disk_usage.go                      | 47 ++++++++++++++---------
 4 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/api/server/router/system/backend.go b/api/server/router/system/backend.go
index 6946c4e2d1..da1de380db 100644
--- a/api/server/router/system/backend.go
+++ b/api/server/router/system/backend.go
@@ -14,7 +14,7 @@ import (
 type Backend interface {
 	SystemInfo() (*types.Info, error)
 	SystemVersion() types.Version
-	SystemDiskUsage() (*types.DiskUsage, error)
+	SystemDiskUsage(ctx context.Context) (*types.DiskUsage, error)
 	SubscribeToEvents(since, until time.Time, ef filters.Args) ([]events.Message, chan interface{})
 	UnsubscribeFromEvents(chan interface{})
 	AuthenticateToRegistry(ctx context.Context, authConfig *types.AuthConfig) (string, string, error)
diff --git a/api/server/router/system/system.go b/api/server/router/system/system.go
index 357bbbe8ee..f8a851cfe3 100644
--- a/api/server/router/system/system.go
+++ b/api/server/router/system/system.go
@@ -26,7 +26,7 @@ func NewRouter(b Backend, c *cluster.Cluster) router.Router {
 		router.NewGetRoute("/events", r.getEvents, router.WithCancel),
 		router.NewGetRoute("/info", r.getInfo),
 		router.NewGetRoute("/version", r.getVersion),
-		router.NewGetRoute("/system/df", r.getDiskUsage),
+		router.NewGetRoute("/system/df", r.getDiskUsage, router.WithCancel),
 		router.NewPostRoute("/auth", r.postAuth),
 	}
 
diff --git a/api/server/router/system/system_routes.go b/api/server/router/system/system_routes.go
index bb7853927d..7d3ba9a6ec 100644
--- a/api/server/router/system/system_routes.go
+++ b/api/server/router/system/system_routes.go
@@ -71,7 +71,7 @@ func (s *systemRouter) getVersion(ctx context.Context, w http.ResponseWriter, r
 }
 
 func (s *systemRouter) getDiskUsage(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
-	du, err := s.backend.SystemDiskUsage()
+	du, err := s.backend.SystemDiskUsage(ctx)
 	if err != nil {
 		return err
 	}
diff --git a/daemon/disk_usage.go b/daemon/disk_usage.go
index fc77a3d8fd..6bff4f0ff9 100644
--- a/daemon/disk_usage.go
+++ b/daemon/disk_usage.go
@@ -3,6 +3,8 @@ package daemon
 import (
 	"fmt"
 
+	"golang.org/x/net/context"
+
 	"github.com/Sirupsen/logrus"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/filters"
@@ -34,7 +36,7 @@ func (daemon *Daemon) getLayerRefs() map[layer.ChainID]int {
 }
 
 // SystemDiskUsage returns information about the daemon data disk usage
-func (daemon *Daemon) SystemDiskUsage() (*types.DiskUsage, error) {
+func (daemon *Daemon) SystemDiskUsage(ctx context.Context) (*types.DiskUsage, error) {
 	// Retrieve container list
 	allContainers, err := daemon.Containers(&types.ContainerListOptions{
 		Size: true,
@@ -53,17 +55,22 @@ func (daemon *Daemon) SystemDiskUsage() (*types.DiskUsage, error) {
 	// Get all local volumes
 	allVolumes := []*types.Volume{}
 	getLocalVols := func(v volume.Volume) error {
-		name := v.Name()
-		refs := daemon.volumes.Refs(v)
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			name := v.Name()
+			refs := daemon.volumes.Refs(v)
 
-		tv := volumeToAPIType(v)
-		sz, err := directory.Size(v.Path())
-		if err != nil {
-			logrus.Warnf("failed to determine size of volume %v", name)
-			sz = -1
+			tv := volumeToAPIType(v)
+			sz, err := directory.Size(v.Path())
+			if err != nil {
+				logrus.Warnf("failed to determine size of volume %v", name)
+				sz = -1
+			}
+			tv.UsageData = &types.VolumeUsageData{Size: sz, RefCount: int64(len(refs))}
+			allVolumes = append(allVolumes, tv)
 		}
-		tv.UsageData = &types.VolumeUsageData{Size: sz, RefCount: int64(len(refs))}
-		allVolumes = append(allVolumes, tv)
 
 		return nil
 	}
@@ -78,17 +85,21 @@ func (daemon *Daemon) SystemDiskUsage() (*types.DiskUsage, error) {
 	allLayers := daemon.layerStore.Map()
 	var allLayersSize int64
 	for _, l := range allLayers {
-		size, err := l.DiffSize()
-		if err == nil {
-			if _, ok := layerRefs[l.ChainID()]; ok {
-				allLayersSize += size
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			size, err := l.DiffSize()
+			if err == nil {
+				if _, ok := layerRefs[l.ChainID()]; ok {
+					allLayersSize += size
+				} else {
+					logrus.Warnf("found leaked image layer %v", l.ChainID())
+				}
 			} else {
-				logrus.Warnf("found leaked image layer %v", l.ChainID())
+				logrus.Warnf("failed to get diff size for layer %v", l.ChainID())
 			}
-		} else {
-			logrus.Warnf("failed to get diff size for layer %v", l.ChainID())
 		}
-
 	}
 
 	return &types.DiskUsage{