From 180f4e8be798a0e8240e40a3a1259db1f728dccc Mon Sep 17 00:00:00 2001
From: Naveen Mahalingam <jedib0t@outlook.com>
Date: Mon, 28 Nov 2022 08:45:08 -0800
Subject: [PATCH] progress: refactor AppendTracker for readability

---
 progress/progress.go     | 21 ++++++---------------
 progress/tracker.go      | 10 ++++++++++
 progress/tracker_test.go | 17 +++++++++++++++++
 3 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/progress/progress.go b/progress/progress.go
index 9d6e4b3..1253ece 100644
--- a/progress/progress.go
+++ b/progress/progress.go
@@ -68,8 +68,9 @@ const (
 // cycle.
 func (p *Progress) AppendTracker(t *Tracker) {
 	t.start()
-
 	p.overallTrackerMutex.Lock()
+	defer p.overallTrackerMutex.Unlock()
+
 	if p.overallTracker == nil {
 		p.overallTracker = &Tracker{Total: 1}
 		if p.numTrackersExpected > 0 {
@@ -78,23 +79,13 @@ func (p *Progress) AppendTracker(t *Tracker) {
 		p.overallTracker.start()
 	}
 
-	// corner case: overall tracker is done, but new tracker is being added. We should recover it.
-	// issue: #245
-	if p.overallTracker.IsDone() {
-		p.overallTracker.mutex.Lock()
-		p.overallTracker.done = false
-		p.overallTracker.mutex.Unlock()
-	}
-
+	// append the tracker to the "in-queue" list
 	p.trackersInQueueMutex.Lock()
 	p.trackersInQueue = append(p.trackersInQueue, t)
 	p.trackersInQueueMutex.Unlock()
-	p.overallTracker.mutex.Lock()
-	if p.overallTracker.Total < int64(p.Length())*100 {
-		p.overallTracker.Total = int64(p.Length()) * 100
-	}
-	p.overallTracker.mutex.Unlock()
-	p.overallTrackerMutex.Unlock()
+
+	// update the expected total progress since we are appending a new tracker
+	p.overallTracker.UpdateTotal(int64(p.Length()) * 100)
 }
 
 // AppendTrackers appends one or more Trackers for tracking.
diff --git a/progress/tracker.go b/progress/tracker.go
index adfb326..6ff16e5 100644
--- a/progress/tracker.go
+++ b/progress/tracker.go
@@ -158,6 +158,16 @@ func (t *Tracker) UpdateMessage(msg string) {
 	t.mutex.Unlock()
 }
 
+// UpdateTotal updates the total value.
+func (t *Tracker) UpdateTotal(total int64) {
+	t.mutex.Lock()
+	if total > t.Total {
+		t.done = false
+	}
+	t.Total = total
+	t.mutex.Unlock()
+}
+
 // Value returns the current value of the tracker.
 func (t *Tracker) Value() int64 {
 	t.mutex.RLock()
diff --git a/progress/tracker_test.go b/progress/tracker_test.go
index 9a48520..dc5d848 100644
--- a/progress/tracker_test.go
+++ b/progress/tracker_test.go
@@ -184,3 +184,20 @@ func TestTracker_UpdateMessage(t *testing.T) {
 	tracker.UpdateMessage("bar")
 	assert.Equal(t, "bar", tracker.message())
 }
+
+func TestTracker_UpdateTotal(t *testing.T) {
+	tracker := Tracker{Total: 100}
+	assert.False(t, tracker.IsDone())
+	tracker.SetValue(100)
+	assert.True(t, tracker.IsDone())
+
+	tracker.UpdateTotal(101)
+	assert.False(t, tracker.IsDone())
+	tracker.SetValue(101)
+	assert.True(t, tracker.IsDone())
+
+	tracker.UpdateTotal(100)
+	assert.True(t, tracker.IsDone())
+	tracker.SetValue(100)
+	assert.True(t, tracker.IsDone())
+}