Support cancellation via context in DiskUsage.

cpuguy83 · cpuguy83 · commit 6cde9049f2f2 · 2018-04-08T13:46:25.000-04:00
Signed-off-by: Brian Goff &lt;cpuguy83@gmail.com&gt;
diff --git a/fs/du.go b/fs/du.go
@@ -10,8 +10,8 @@ type Usage struct {
 
 // DiskUsage counts the number of inodes and disk usage for the resources under
 // path.
-func DiskUsage(roots ...string) (Usage, error) {
-	return diskUsage(roots...)
+func DiskUsage(ctx context.Context, roots ...string) (Usage, error) {
+	return diskUsage(ctx, roots...)
 }
 
 // DiffUsage counts the numbers of inodes and disk usage in the
diff --git a/fs/du_unix.go b/fs/du_unix.go
@@ -24,7 +24,7 @@ func newInode(stat *syscall.Stat_t) inode {
 	}
 }
 
-func diskUsage(roots ...string) (Usage, error) {
+func diskUsage(ctx context.Context, roots ...string) (Usage, error) {
 
 	var (
 		size   int64
@@ -37,6 +37,12 @@ func diskUsage(roots ...string) (Usage, error) {
 				return err
 			}
 
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+			}
+
 			inoKey := newInode(fi.Sys().(*syscall.Stat_t))
 			if _, ok := inodes[inoKey]; !ok {
 				inodes[inoKey] = struct{}{}
diff --git a/fs/du_windows.go b/fs/du_windows.go
@@ -8,7 +8,7 @@ import (
 	"path/filepath"
 )
 
-func diskUsage(roots ...string) (Usage, error) {
+func diskUsage(ctx context.Context, roots ...string) (Usage, error) {
 	var (
 		size int64
 	)
@@ -21,6 +21,12 @@ func diskUsage(roots ...string) (Usage, error) {
 				return err
 			}
 
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+			}
+
 			size += fi.Size()
 			return nil
 		}); err != nil {

Original file line number	Diff line number	Diff line change
`@@ -10,8 +10,8 @@ type Usage struct {`
`10`	`10`
`11`	`11`	`// DiskUsage counts the number of inodes and disk usage for the resources under`
`12`	`12`	`// path.`
`13`		`-func DiskUsage(roots ...string) (Usage, error) {`
`14`		`- return diskUsage(roots...)`
	`13`	`+func DiskUsage(ctx context.Context, roots ...string) (Usage, error) {`
	`14`	`+ return diskUsage(ctx, roots...)`
`15`	`15`	`}`
`16`	`16`
`17`	`17`	`// DiffUsage counts the numbers of inodes and disk usage in the`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ func newInode(stat *syscall.Stat_t) inode {`
`24`	`24`	`}`
`25`	`25`	`}`
`26`	`26`
`27`		`-func diskUsage(roots ...string) (Usage, error) {`
	`27`	`+func diskUsage(ctx context.Context, roots ...string) (Usage, error) {`
`28`	`28`
`29`	`29`	`var (`
`30`	`30`	`size int64`
`@@ -37,6 +37,12 @@ func diskUsage(roots ...string) (Usage, error) {`
`37`	`37`	`return err`
`38`	`38`	`}`
`39`	`39`
	`40`	`+ select {`
	`41`	`+ case <-ctx.Done():`
	`42`	`+ return ctx.Err()`
	`43`	`+ default:`
	`44`	`+ }`
	`45`	`+`
`40`	`46`	`inoKey := newInode(fi.Sys().(*syscall.Stat_t))`
`41`	`47`	`if _, ok := inodes[inoKey]; !ok {`
`42`	`48`	`inodes[inoKey] = struct{}{}`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ import (`
`8`	`8`	`"path/filepath"`
`9`	`9`	`)`
`10`	`10`
`11`		`-func diskUsage(roots ...string) (Usage, error) {`
	`11`	`+func diskUsage(ctx context.Context, roots ...string) (Usage, error) {`
`12`	`12`	`var (`
`13`	`13`	`size int64`
`14`	`14`	`)`
`@@ -21,6 +21,12 @@ func diskUsage(roots ...string) (Usage, error) {`
`21`	`21`	`return err`
`22`	`22`	`}`
`23`	`23`
	`24`	`+ select {`
	`25`	`+ case <-ctx.Done():`
	`26`	`+ return ctx.Err()`
	`27`	`+ default:`
	`28`	`+ }`
	`29`	`+`
`24`	`30`	`size += fi.Size()`
`25`	`31`	`return nil`
`26`	`32`	`}); err != nil {`