From cf1573f9363107e30399011565bf11f95bac52d7 Mon Sep 17 00:00:00 2001
From: Chris Suszynski <ksuszyns@redhat.com>
Date: Fri, 22 Mar 2024 18:17:49 +0100
Subject: [PATCH] :lady_beetle: Fix the hang of the spinner and progress bar on
 stop (#162)

* Fix the hang of the spinner on stop

* Fix the ProgressBar hanging forever at quit

* Wait for natural end for progress bar
---
 pkg/output/tui/progress.go | 25 ++++++++++++++++---------
 pkg/output/tui/spinner.go  | 12 +++++++++---
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/pkg/output/tui/progress.go b/pkg/output/tui/progress.go
index b9108a38..cb26f70f 100644
--- a/pkg/output/tui/progress.go
+++ b/pkg/output/tui/progress.go
@@ -78,7 +78,7 @@ type BubbleProgress struct {
 	speed      int
 	prevSpeed  []int
 	err        error
-	ended      chan struct{}
+	quitChan   chan struct{}
 }
 
 func (b *BubbleProgress) With(fn func(ProgressControl) error) error {
@@ -175,17 +175,23 @@ func (b bubbleProgressHandler) speedChange() (tea.Model, tea.Cmd) {
 }
 
 func (b bubbleProgressHandler) percentChange(event percentChange) (tea.Model, tea.Cmd) {
-	var cmds []tea.Cmd
-
+	cmds := make([]tea.Cmd, 0, 1)
 	cmds = append(cmds, b.prog.SetPercent(float64(event)))
 
 	if event >= 1.0 {
-		cmds = append(cmds, tea.Sequence(b.finalPause(), tea.Quit))
+		cmds = append(cmds, b.quitSignal())
 	}
 
 	return b, tea.Batch(cmds...)
 }
 
+func (b *BubbleProgress) quitSignal() tea.Cmd {
+	// The final pause is to give the progress bar a chance to finish its
+	// animation before quitting. Otherwise, it ends abruptly, and the user
+	// might not see the progress bar at 100%.
+	return tea.Sequence(b.finalPause(), tea.Quit)
+}
+
 func (b bubbleProgressHandler) progressFrame(event progress.FrameMsg) (tea.Model, tea.Cmd) {
 	progressModel, cmd := b.prog.Update(event)
 	if m, ok := progressModel.(progress.Model); ok {
@@ -235,11 +241,11 @@ func (b *BubbleProgress) start() {
 		tea.WithInput(b.InOrStdin()),
 		tea.WithOutput(out),
 	)
-	b.ended = make(chan struct{})
+	b.quitChan = make(chan struct{})
 	go func() {
 		t := b.tea
 		_, _ = t.Run()
-		close(b.ended)
+		close(b.quitChan)
 		if term.IsTerminal(out) {
 			if err := t.ReleaseTerminal(); err != nil {
 				panic(err)
@@ -252,11 +258,12 @@ func (b *BubbleProgress) stop() {
 	if b.tea == nil {
 		return
 	}
-	b.tea.Wait()
 
-	<-b.ended
+	b.tea.Send(b.quitSignal())
+	<-b.quitChan
+
 	b.tea = nil
-	b.ended = nil
+	b.quitChan = nil
 }
 
 func (b *BubbleProgress) onProgress(percent float64) {
diff --git a/pkg/output/tui/spinner.go b/pkg/output/tui/spinner.go
index 91fd762e..438ef358 100644
--- a/pkg/output/tui/spinner.go
+++ b/pkg/output/tui/spinner.go
@@ -43,8 +43,9 @@ type BubbleSpinner struct {
 	output.InputOutput
 	Message
 
-	spin spinner.Model
-	tea  *tea.Program
+	spin     spinner.Model
+	tea      *tea.Program
+	quitChan chan struct{}
 }
 
 func (b *BubbleSpinner) With(fn func(Spinner) error) error {
@@ -77,9 +78,11 @@ func (b *BubbleSpinner) start() {
 		tea.WithInput(b.InOrStdin()),
 		tea.WithOutput(out),
 	)
+	b.quitChan = make(chan struct{})
 	go func() {
 		t := b.tea
 		_, _ = t.Run()
+		close(b.quitChan)
 		if term.IsTerminal(out) {
 			_ = t.ReleaseTerminal()
 		}
@@ -90,9 +93,12 @@ func (b *BubbleSpinner) stop() {
 	if b.tea == nil {
 		return
 	}
+
 	b.tea.Quit()
-	b.tea.Wait()
+	<-b.quitChan
+
 	b.tea = nil
+	b.quitChan = nil
 	endMsg := fmt.Sprintf("%s %s\n",
 		b.Message.Text, spinnerStyle().Render("Done"))
 	_, _ = b.OutOrStdout().Write([]byte(endMsg))