From 505f8a22298ab0c1fe93aa1a1ce475e340d24228 Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Sat, 5 Dec 2020 23:57:06 +0100
Subject: [PATCH] progress/counter: Support updating the progress bar maximum

---
 cmd/restic/progress.go               | 13 ++++++++-----
 internal/ui/progress/counter.go      | 29 ++++++++++++++++++++++++----
 internal/ui/progress/counter_test.go | 18 +++++++++++++----
 3 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/cmd/restic/progress.go b/cmd/restic/progress.go
index 2fbd97c6c..c0b6c56fb 100644
--- a/cmd/restic/progress.go
+++ b/cmd/restic/progress.go
@@ -33,11 +33,14 @@ func newProgressMax(show bool, max uint64, description string) *progress.Counter
 	}
 	interval := calculateProgressInterval()
 
-	return progress.New(interval, func(v uint64, d time.Duration, final bool) {
-		status := fmt.Sprintf("[%s] %s  %d / %d %s",
-			formatDuration(d),
-			formatPercent(v, max),
-			v, max, description)
+	return progress.New(interval, max, func(v uint64, max uint64, d time.Duration, final bool) {
+		var status string
+		if max == 0 {
+			status = fmt.Sprintf("[%s]          %d %s", formatDuration(d), v, description)
+		} else {
+			status = fmt.Sprintf("[%s] %s  %d / %d %s",
+				formatDuration(d), formatPercent(v, max), v, max, description)
+		}
 
 		if w := stdoutTerminalWidth(); w > 0 {
 			status = shortenStatus(w, status)
diff --git a/internal/ui/progress/counter.go b/internal/ui/progress/counter.go
index bf4906978..82d5e39ba 100644
--- a/internal/ui/progress/counter.go
+++ b/internal/ui/progress/counter.go
@@ -12,7 +12,7 @@ import (
 //
 // The final argument is true if Counter.Done has been called,
 // which means that the current call will be the last.
-type Func func(value uint64, runtime time.Duration, final bool)
+type Func func(value uint64, total uint64, runtime time.Duration, final bool)
 
 // A Counter tracks a running count and controls a goroutine that passes its
 // value periodically to a Func.
@@ -27,16 +27,19 @@ type Counter struct {
 
 	valueMutex sync.Mutex
 	value      uint64
+	max        uint64
 }
 
 // New starts a new Counter.
-func New(interval time.Duration, report Func) *Counter {
+func New(interval time.Duration, total uint64, report Func) *Counter {
 	c := &Counter{
 		report:  report,
 		start:   time.Now(),
 		stopped: make(chan struct{}),
 		stop:    make(chan struct{}),
+		max:     total,
 	}
+
 	if interval > 0 {
 		c.tick = time.NewTicker(interval)
 	}
@@ -56,6 +59,16 @@ func (c *Counter) Add(v uint64) {
 	c.valueMutex.Unlock()
 }
 
+// SetMax sets the maximum expected counter value. This method is concurrency-safe.
+func (c *Counter) SetMax(max uint64) {
+	if c == nil {
+		return
+	}
+	c.valueMutex.Lock()
+	c.max = max
+	c.valueMutex.Unlock()
+}
+
 // Done tells a Counter to stop and waits for it to report its final value.
 func (c *Counter) Done() {
 	if c == nil {
@@ -77,11 +90,19 @@ func (c *Counter) get() uint64 {
 	return v
 }
 
+func (c *Counter) getMax() uint64 {
+	c.valueMutex.Lock()
+	max := c.max
+	c.valueMutex.Unlock()
+
+	return max
+}
+
 func (c *Counter) run() {
 	defer close(c.stopped)
 	defer func() {
 		// Must be a func so that time.Since isn't called at defer time.
-		c.report(c.get(), time.Since(c.start), true)
+		c.report(c.get(), c.getMax(), time.Since(c.start), true)
 	}()
 
 	var tick <-chan time.Time
@@ -101,6 +122,6 @@ func (c *Counter) run() {
 			return
 		}
 
-		c.report(c.get(), now.Sub(c.start), false)
+		c.report(c.get(), c.getMax(), now.Sub(c.start), false)
 	}
 }
diff --git a/internal/ui/progress/counter_test.go b/internal/ui/progress/counter_test.go
index 9a76d9cbf..84d1e5a64 100644
--- a/internal/ui/progress/counter_test.go
+++ b/internal/ui/progress/counter_test.go
@@ -10,23 +10,30 @@ import (
 
 func TestCounter(t *testing.T) {
 	const N = 100
+	const startTotal = uint64(12345)
 
 	var (
 		finalSeen  = false
 		increasing = true
 		last       uint64
+		lastTotal  = startTotal
 		ncalls     int
+		nmaxChange int
 	)
 
-	report := func(value uint64, d time.Duration, final bool) {
+	report := func(value uint64, total uint64, d time.Duration, final bool) {
 		finalSeen = true
 		if value < last {
 			increasing = false
 		}
 		last = value
+		if total != lastTotal {
+			nmaxChange++
+		}
+		lastTotal = total
 		ncalls++
 	}
-	c := progress.New(10*time.Millisecond, report)
+	c := progress.New(10*time.Millisecond, startTotal, report)
 
 	done := make(chan struct{})
 	go func() {
@@ -35,6 +42,7 @@ func TestCounter(t *testing.T) {
 			time.Sleep(time.Millisecond)
 			c.Add(1)
 		}
+		c.SetMax(42)
 	}()
 
 	<-done
@@ -43,6 +51,8 @@ func TestCounter(t *testing.T) {
 	test.Assert(t, finalSeen, "final call did not happen")
 	test.Assert(t, increasing, "values not increasing")
 	test.Equals(t, uint64(N), last)
+	test.Equals(t, uint64(42), lastTotal)
+	test.Equals(t, int(1), nmaxChange)
 
 	t.Log("number of calls:", ncalls)
 }
@@ -58,14 +68,14 @@ func TestCounterNoTick(t *testing.T) {
 	finalSeen := false
 	otherSeen := false
 
-	report := func(value uint64, d time.Duration, final bool) {
+	report := func(value, total uint64, d time.Duration, final bool) {
 		if final {
 			finalSeen = true
 		} else {
 			otherSeen = true
 		}
 	}
-	c := progress.New(0, report)
+	c := progress.New(0, 1, report)
 	time.Sleep(time.Millisecond)
 	c.Done()