Merge pull request #4157 from cpuguy83/23.0_improve_cp_progress

[23.0 backport] improve cp progress

cli/command/container/cp.go

@@ -1,15 +1,20 @@
 package container
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
+	"time"
 
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
+	"github.com/docker/cli/cli/streams"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/pkg/archive"
 	"github.com/docker/docker/pkg/system"
@@ -48,28 +53,75 @@ type cpConfig struct {
 // copying files to/from a container.
 type copyProgressPrinter struct {
 	io.ReadCloser
-	toContainer bool
-	total       *float64
-	writer      io.Writer
+	total *int64
 }
 
+const (
+	copyToContainerHeader       = "Copying to container - "
+	copyFromContainerHeader     = "Copying from container - "
+	copyProgressUpdateThreshold = 75 * time.Millisecond
+)
+
 func (pt *copyProgressPrinter) Read(p []byte) (int, error) {
 	n, err := pt.ReadCloser.Read(p)
-	*pt.total += float64(n)
-
-	if err == nil {
-		fmt.Fprint(pt.writer, aec.Restore)
-		fmt.Fprint(pt.writer, aec.EraseLine(aec.EraseModes.All))
-		if pt.toContainer {
-			fmt.Fprintln(pt.writer, "Copying to container - "+units.HumanSize(*pt.total))
-		} else {
-			fmt.Fprintln(pt.writer, "Copying from container - "+units.HumanSize(*pt.total))
-		}
-	}
-
+	atomic.AddInt64(pt.total, int64(n))
 	return n, err
 }
 
+func copyProgress(ctx context.Context, dst io.Writer, header string, total *int64) (func(), <-chan struct{}) {
+	done := make(chan struct{})
+	if !streams.NewOut(dst).IsTerminal() {
+		close(done)
+		return func() {}, done
+	}
+
+	fmt.Fprint(dst, aec.Save)
+	fmt.Fprint(dst, "Preparing to copy...")
+
+	restore := func() {
+		fmt.Fprint(dst, aec.Restore)
+		fmt.Fprint(dst, aec.EraseLine(aec.EraseModes.All))
+	}
+
+	go func() {
+		defer close(done)
+		fmt.Fprint(dst, aec.Hide)
+		defer fmt.Fprint(dst, aec.Show)
+
+		fmt.Fprint(dst, aec.Restore)
+		fmt.Fprint(dst, aec.EraseLine(aec.EraseModes.All))
+		fmt.Fprint(dst, header)
+
+		var last int64
+		fmt.Fprint(dst, progressHumanSize(last))
+
+		buf := bytes.NewBuffer(nil)
+		ticker := time.NewTicker(copyProgressUpdateThreshold)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				n := atomic.LoadInt64(total)
+				if n == last {
+					// Don't write to the terminal, if we don't need to.
+					continue
+				}
+
+				// Write to the buffer first to avoid flickering and context switching
+				fmt.Fprint(buf, aec.Column(uint(len(header)+1)))
+				fmt.Fprint(buf, aec.EraseLine(aec.EraseModes.Tail))
+				fmt.Fprint(buf, progressHumanSize(n))
+
+				buf.WriteTo(dst)
+				buf.Reset()
+				last += n
+			}
+		}
+	}()
+	return restore, done
+}
+
 // NewCopyCommand creates a new `docker cp` command
 func NewCopyCommand(dockerCli command.Cli) *cobra.Command {
 	var opts copyOptions
@@ -113,6 +165,10 @@ func NewCopyCommand(dockerCli command.Cli) *cobra.Command {
 	return cmd
 }
 
+func progressHumanSize(n int64) string {
+	return units.HumanSizeWithPrecision(float64(n), 3)
+}
+
 func runCopy(dockerCli command.Cli, opts copyOptions) error {
 	srcContainer, srcPath := splitCpArg(opts.source)
 	destContainer, destPath := splitCpArg(opts.destination)
@@ -193,6 +249,9 @@ func copyFromContainer(ctx context.Context, dockerCli command.Cli, copyConfig cp
 
 	}
 
+	ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
+	defer cancel()
+
 	content, stat, err := client.CopyFromContainer(ctx, copyConfig.container, srcPath)
 	if err != nil {
 		return err
@@ -211,12 +270,10 @@ func copyFromContainer(ctx context.Context, dockerCli command.Cli, copyConfig cp
 		RebaseName: rebaseName,
 	}
 
-	var copiedSize float64
+	var copiedSize int64
 	if !copyConfig.quiet {
 		content = &copyProgressPrinter{
 			ReadCloser: content,
-			toContainer: false,
-			writer:      dockerCli.Err(),
 			total:      &copiedSize,
 		}
 	}
@@ -231,12 +288,12 @@ func copyFromContainer(ctx context.Context, dockerCli command.Cli, copyConfig cp
 		return archive.CopyTo(preArchive, srcInfo, dstPath)
 	}
 
-	fmt.Fprint(dockerCli.Err(), aec.Save)
-	fmt.Fprintln(dockerCli.Err(), "Preparing to copy...")
+	restore, done := copyProgress(ctx, dockerCli.Err(), copyFromContainerHeader, &copiedSize)
 	res := archive.CopyTo(preArchive, srcInfo, dstPath)
-	fmt.Fprint(dockerCli.Err(), aec.Restore)
-	fmt.Fprint(dockerCli.Err(), aec.EraseLine(aec.EraseModes.All))
-	fmt.Fprintln(dockerCli.Err(), "Successfully copied", units.HumanSize(copiedSize), "to", dstPath)
+	cancel()
+	<-done
+	restore()
+	fmt.Fprintln(dockerCli.Err(), "Successfully copied", progressHumanSize(copiedSize), "to", dstPath)
 
 	return res
 }
@@ -293,7 +350,7 @@ func copyToContainer(ctx context.Context, dockerCli command.Cli, copyConfig cpCo
 	var (
 		content         io.ReadCloser
 		resolvedDstPath string
-		copiedSize      float64
+		copiedSize      int64
 	)
 
 	if srcPath == "-" {
@@ -338,8 +395,6 @@ func copyToContainer(ctx context.Context, dockerCli command.Cli, copyConfig cpCo
 		if !copyConfig.quiet {
 			content = &copyProgressPrinter{
 				ReadCloser: content,
-				toContainer: true,
-				writer:      dockerCli.Err(),
 				total:      &copiedSize,
 			}
 		}
@@ -354,12 +409,13 @@ func copyToContainer(ctx context.Context, dockerCli command.Cli, copyConfig cpCo
 		return client.CopyToContainer(ctx, copyConfig.container, resolvedDstPath, content, options)
 	}
 
-	fmt.Fprint(dockerCli.Err(), aec.Save)
-	fmt.Fprintln(dockerCli.Err(), "Preparing to copy...")
+	ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
+	restore, done := copyProgress(ctx, dockerCli.Err(), copyToContainerHeader, &copiedSize)
 	res := client.CopyToContainer(ctx, copyConfig.container, resolvedDstPath, content, options)
-	fmt.Fprint(dockerCli.Err(), aec.Restore)
-	fmt.Fprint(dockerCli.Err(), aec.EraseLine(aec.EraseModes.All))
-	fmt.Fprintln(dockerCli.Err(), "Successfully copied", units.HumanSize(copiedSize), "to", copyConfig.container+":"+dstInfo.Path)
+	cancel()
+	<-done
+	restore()
+	fmt.Fprintln(dockerCli.Err(), "Successfully copied", progressHumanSize(copiedSize), "to", copyConfig.container+":"+dstInfo.Path)
 
 	return res
 }