From d52f5db6ba450a879ee190973d0126cb632af86e Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Mon, 6 Apr 2020 19:03:29 -0700
Subject: [PATCH] commands: add du command

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 commands/diskusage.go | 137 ++++++++++++++++++++++++++++++++++++------
 commands/root.go      |   1 +
 go.mod                |   1 +
 3 files changed, 120 insertions(+), 19 deletions(-)

diff --git a/commands/diskusage.go b/commands/diskusage.go
index bdd3597c..3ae434b7 100644
--- a/commands/diskusage.go
+++ b/commands/diskusage.go
@@ -6,10 +6,117 @@ import (
 	"os"
 	"text/tabwriter"
 
+	"github.com/docker/buildx/build"
+	"github.com/docker/cli/cli"
+	"github.com/docker/cli/cli/command"
+	"github.com/docker/cli/opts"
 	"github.com/moby/buildkit/client"
+	"github.com/moby/buildkit/util/appcontext"
+	"github.com/spf13/cobra"
 	"github.com/tonistiigi/units"
+	"golang.org/x/sync/errgroup"
 )
 
+type duOptions struct {
+	filter  opts.FilterOpt
+	verbose bool
+}
+
+func runDiskUsage(dockerCli command.Cli, opts duOptions) error {
+	ctx := appcontext.Context()
+
+	pi, err := toBuildkitPruneInfo(opts.filter.Value())
+	if err != nil {
+		return err
+	}
+
+	dis, err := getDefaultDrivers(ctx, dockerCli, "")
+	if err != nil {
+		return err
+	}
+
+	for _, di := range dis {
+		if di.Err != nil {
+			return err
+		}
+	}
+
+	out := make([][]*client.UsageInfo, len(dis))
+
+	eg, ctx := errgroup.WithContext(ctx)
+	for i, di := range dis {
+		func(i int, di build.DriverInfo) {
+			eg.Go(func() error {
+				if di.Driver != nil {
+					c, err := di.Driver.Client(ctx)
+					if err != nil {
+						return err
+					}
+					du, err := c.DiskUsage(ctx, client.WithFilter(pi.Filter))
+					if err != nil {
+						return err
+					}
+					out[i] = du
+					return nil
+				}
+				return nil
+			})
+		}(i, di)
+	}
+
+	if err := eg.Wait(); err != nil {
+		return err
+	}
+
+	tw := tabwriter.NewWriter(os.Stdout, 1, 8, 1, '\t', 0)
+	first := true
+	for _, du := range out {
+		if du == nil {
+			continue
+		}
+		if opts.verbose {
+			printVerbose(tw, du)
+		} else {
+			if first {
+				printTableHeader(tw)
+				first = false
+			}
+			for _, di := range du {
+				printTableRow(tw, di)
+			}
+
+			tw.Flush()
+		}
+	}
+
+	if opts.filter.Value().Len() == 0 {
+		printSummary(tw, out)
+	}
+
+	tw.Flush()
+	return nil
+}
+
+func duCmd(dockerCli command.Cli) *cobra.Command {
+	options := duOptions{filter: opts.NewFilterOpt()}
+
+	cmd := &cobra.Command{
+		Use:   "du",
+		Short: "Disk usage",
+		Args:  cli.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return runDiskUsage(dockerCli, options)
+		},
+		Annotations: map[string]string{"version": "1.00"},
+	}
+
+	flags := cmd.Flags()
+	flags.Var(&options.filter, "filter", "Provide filter values")
+	flags.BoolVar(&options.verbose, "verbose", false, "Provide a more verbose output")
+
+	return cmd
+}
+
 func printKV(w io.Writer, k string, v interface{}) {
 	fmt.Fprintf(w, "%s:\t%v\n", k, v)
 }
@@ -42,16 +149,6 @@ func printVerbose(tw *tabwriter.Writer, du []*client.UsageInfo) {
 	tw.Flush()
 }
 
-func printTable(tw *tabwriter.Writer, du []*client.UsageInfo) {
-	printTableHeader(tw)
-
-	for _, di := range du {
-		printTableRow(tw, di)
-	}
-
-	tw.Flush()
-}
-
 func printTableHeader(tw *tabwriter.Writer) {
 	fmt.Fprintln(tw, "ID\tRECLAIMABLE\tSIZE\tLAST ACCESSED")
 }
@@ -68,20 +165,22 @@ func printTableRow(tw *tabwriter.Writer, di *client.UsageInfo) {
 	fmt.Fprintf(tw, "%-71s\t%-11v\t%s\t\n", id, !di.InUse, size)
 }
 
-func printSummary(tw *tabwriter.Writer, du []*client.UsageInfo) {
+func printSummary(tw *tabwriter.Writer, dus [][]*client.UsageInfo) {
 	total := int64(0)
 	reclaimable := int64(0)
 	shared := int64(0)
 
-	for _, di := range du {
-		if di.Size > 0 {
-			total += di.Size
-			if !di.InUse {
-				reclaimable += di.Size
+	for _, du := range dus {
+		for _, di := range du {
+			if di.Size > 0 {
+				total += di.Size
+				if !di.InUse {
+					reclaimable += di.Size
+				}
+			}
+			if di.Shared {
+				shared += di.Size
 			}
-		}
-		if di.Shared {
-			shared += di.Size
 		}
 	}
 
diff --git a/commands/root.go b/commands/root.go
index 2adfb5dc..0f141e0e 100644
--- a/commands/root.go
+++ b/commands/root.go
@@ -36,6 +36,7 @@ func addCommands(cmd *cobra.Command, dockerCli command.Cli) {
 		uninstallCmd(dockerCli),
 		versionCmd(dockerCli),
 		pruneCmd(dockerCli),
+		duCmd(dockerCli),
 		imagetoolscmd.RootCmd(dockerCli),
 	)
 }
diff --git a/go.mod b/go.mod
index 7616d6e2..8f71b06c 100644
--- a/go.mod
+++ b/go.mod
@@ -53,6 +53,7 @@ require (
 	github.com/spf13/viper v1.3.2 // indirect
 	github.com/stretchr/testify v1.4.0
 	github.com/theupdateframework/notary v0.6.1 // indirect
+	github.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea
 	github.com/xlab/handysort v0.0.0-20150421192137-fb3537ed64a1 // indirect
 	golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e
 	gopkg.in/dancannon/gorethink.v3 v3.0.5 // indirect