progress: allow printing after progress is done

When something needs to be printed after the progress bar is done, it
was not possible to do so without messing up the terminal. This patch
adds a mutex around the buffer used to store output while the progress
bar is active. This allows to safely print after stopping the progress
bar.

Author: lzap

pkg/progress/progress.go

@@ -1,12 +1,14 @@
 package progress
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
 	"log"
 	"os"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/cheggaaa/pb/v3"
@@ -70,6 +72,11 @@ type ProgressBar interface {
 	// Stop will stop rendering the progress information, the
 	// screen is not cleared, the last few lines will be visible
 	Stop()
+
+	// Write implements io.Writer that is used to capture output
+	// that is written after the progress bar is stopped via
+	// Stop().
+	Write(p []byte) (n int, err error)
 }
 
 // New creates a new progressbar based on the requested type
@@ -103,7 +110,9 @@ type terminalProgressBar struct {
 
 	shutdownCh chan bool
 
-	out io.Writer
+	out   io.Writer
+	buf   bytes.Buffer
+	bufMu sync.Mutex
 }
 
 // NewTerminalProgressBar creates a new default pb3 based progressbar suitable for
@@ -219,6 +228,13 @@ func (b *terminalProgressBar) Err() error {
 	return errors.Join(errs...)
 }
 
+func (b *terminalProgressBar) Write(p []byte) (n int, err error) {
+	b.bufMu.Lock()
+	defer b.bufMu.Unlock()
+
+	return b.buf.Write(p)
+}
+
 func (b *terminalProgressBar) Stop() {
 	if b.shutdownCh == nil {
 		return
@@ -242,6 +258,13 @@ func (b *terminalProgressBar) Stop() {
 	if err := b.Err(); err != nil {
 		fmt.Fprintf(b.out, "error from pb.ProgressBar: %v", err)
 	}
+	// write any buffered output
+	if b.buf.Len() > 0 {
+		b.bufMu.Lock()
+		defer b.bufMu.Unlock()
+		fmt.Fprintf(b.out, "%s%s", ERASE_LINE, b.buf.String())
+		b.buf.Reset()
+	}
 }
 
 type verboseProgressBar struct {
@@ -268,6 +291,10 @@ func (b *verboseProgressBar) SetMessagef(msg string, args ...interface{}) {
 func (b *verboseProgressBar) Start() {
 }
 
+func (b *verboseProgressBar) Write(p []byte) (n int, err error) {
+	return b.w.Write(p)
+}
+
 func (b *verboseProgressBar) Stop() {
 }
 
@@ -304,6 +331,13 @@ func (b *debugProgressBar) Start() {
 	fmt.Fprintf(b.w, "Start progressbar\n")
 }
 
+func (b *debugProgressBar) Write(p []byte) (n int, err error) {
+	fmt.Fprintf(b.w, "write: ")
+	n, err = b.w.Write(p)
+	fmt.Fprintf(b.w, "\n")
+	return
+}
+
 func (b *debugProgressBar) Stop() {
 	fmt.Fprintf(b.w, "Stop progressbar\n")
 }