From 7f0aa49f45742baaff195587cc3e4f820c6f42db Mon Sep 17 00:00:00 2001
From: greatroar <61184462+greatroar@users.noreply.github.com>
Date: Sun, 29 Aug 2021 14:55:33 +0200
Subject: [PATCH] cmd/restic: Streamline progress printing

* PrintProgress no longer does unnecessary Sprintf calls, and performs
  fewer allocations in general
* newProgressMax's callback checks whether the terminal supports
  line updates once instead of once per call
* the callback looks up the terminal width once per call instead of
  twice (on Windows)
* the status shortening now uses the Unicode-aware version from
  internal/ui/termstatus (future-proofing)
---
 cmd/restic/cleanup.go   |  2 +-
 cmd/restic/cmd_prune.go | 12 -----------
 cmd/restic/global.go    | 48 ++++++++++++-----------------------------
 cmd/restic/progress.go  | 36 ++++++++++++++++++++++++++-----
 4 files changed, 46 insertions(+), 52 deletions(-)

diff --git a/cmd/restic/cleanup.go b/cmd/restic/cleanup.go
index 738ec590b..67a007d59 100644
--- a/cmd/restic/cleanup.go
+++ b/cmd/restic/cleanup.go
@@ -58,7 +58,7 @@ func RunCleanupHandlers() {
 func CleanupHandler(c <-chan os.Signal) {
 	for s := range c {
 		debug.Log("signal %v received, cleaning up", s)
-		Warnf("%ssignal %v received, cleaning up\n", ClearLine(), s)
+		Warnf("%ssignal %v received, cleaning up\n", clearLine(0), s)
 
 		code := 0
 
diff --git a/cmd/restic/cmd_prune.go b/cmd/restic/cmd_prune.go
index e944c686a..d621afdad 100644
--- a/cmd/restic/cmd_prune.go
+++ b/cmd/restic/cmd_prune.go
@@ -119,18 +119,6 @@ func verifyPruneOptions(opts *PruneOptions) error {
 	return nil
 }
 
-func shortenStatus(maxLength int, s string) string {
-	if len(s) <= maxLength {
-		return s
-	}
-
-	if maxLength < 3 {
-		return s[:maxLength]
-	}
-
-	return s[:maxLength-3] + "..."
-}
-
 func runPrune(opts PruneOptions, gopts GlobalOptions) error {
 	err := verifyPruneOptions(&opts)
 	if err != nil {
diff --git a/cmd/restic/global.go b/cmd/restic/global.go
index 42da399d0..70ec84058 100644
--- a/cmd/restic/global.go
+++ b/cmd/restic/global.go
@@ -197,16 +197,21 @@ func restoreTerminal() {
 }
 
 // ClearLine creates a platform dependent string to clear the current
-// line, so it can be overwritten. ANSI sequences are not supported on
-// current windows cmd shell.
-func ClearLine() string {
-	if runtime.GOOS == "windows" {
-		if w := stdoutTerminalWidth(); w > 0 {
-			return strings.Repeat(" ", w-1) + "\r"
-		}
-		return ""
+// line, so it can be overwritten.
+//
+// w should be the terminal width, or 0 to let clearLine figure it out.
+func clearLine(w int) string {
+	if runtime.GOOS != "windows" {
+		return "\x1b[2K"
 	}
-	return "\x1b[2K"
+
+	// ANSI sequences are not supported on Windows cmd shell.
+	if w <= 0 {
+		if w = stdoutTerminalWidth(); w <= 0 {
+			return ""
+		}
+	}
+	return strings.Repeat(" ", w-1) + "\r"
 }
 
 // Printf writes the message to the configured stdout stream.
@@ -247,31 +252,6 @@ func Verboseff(format string, args ...interface{}) {
 	}
 }
 
-// PrintProgress wraps fmt.Printf to handle the difference in writing progress
-// information to terminals and non-terminal stdout
-func PrintProgress(format string, args ...interface{}) {
-	var (
-		message         string
-		carriageControl string
-	)
-	message = fmt.Sprintf(format, args...)
-
-	if !(strings.HasSuffix(message, "\r") || strings.HasSuffix(message, "\n")) {
-		if stdoutCanUpdateStatus() {
-			carriageControl = "\r"
-		} else {
-			carriageControl = "\n"
-		}
-		message = fmt.Sprintf("%s%s", message, carriageControl)
-	}
-
-	if stdoutCanUpdateStatus() {
-		message = fmt.Sprintf("%s%s", ClearLine(), message)
-	}
-
-	fmt.Print(message)
-}
-
 // Warnf writes the message to the configured stderr stream.
 func Warnf(format string, args ...interface{}) {
 	_, err := fmt.Fprintf(globalOptions.stderr, format, args...)
diff --git a/cmd/restic/progress.go b/cmd/restic/progress.go
index 0c2a24271..26a694c25 100644
--- a/cmd/restic/progress.go
+++ b/cmd/restic/progress.go
@@ -4,9 +4,11 @@ import (
 	"fmt"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/restic/restic/internal/ui/progress"
+	"github.com/restic/restic/internal/ui/termstatus"
 )
 
 // calculateProgressInterval returns the interval configured via RESTIC_PROGRESS_FPS
@@ -32,6 +34,7 @@ func newProgressMax(show bool, max uint64, description string) *progress.Counter
 		return nil
 	}
 	interval := calculateProgressInterval(show)
+	canUpdateStatus := stdoutCanUpdateStatus()
 
 	return progress.New(interval, max, func(v uint64, max uint64, d time.Duration, final bool) {
 		var status string
@@ -42,13 +45,36 @@ func newProgressMax(show bool, max uint64, description string) *progress.Counter
 				formatDuration(d), formatPercent(v, max), v, max, description)
 		}
 
-		if w := stdoutTerminalWidth(); w > 0 {
-			status = shortenStatus(w, status)
-		}
-
-		PrintProgress("%s", status)
+		printProgress(status, canUpdateStatus)
 		if final {
 			fmt.Print("\n")
 		}
 	})
 }
+
+func printProgress(status string, canUpdateStatus bool) {
+	w := stdoutTerminalWidth()
+	if w > 0 {
+		if w < 3 {
+			status = termstatus.Truncate(status, w)
+		} else {
+			status = termstatus.Truncate(status, w-3) + "..."
+		}
+	}
+
+	var carriageControl, clear string
+
+	if canUpdateStatus {
+		clear = clearLine(w)
+	}
+
+	if !(strings.HasSuffix(status, "\r") || strings.HasSuffix(status, "\n")) {
+		if canUpdateStatus {
+			carriageControl = "\r"
+		} else {
+			carriageControl = "\n"
+		}
+	}
+
+	_, _ = os.Stdout.Write([]byte(clear + status + carriageControl))
+}