fix: fatal errors break the terminal due to incorrect bubble tea shutdown (#422)

Author: marktheunissen

cli/benchmark.go

@@ -126,6 +126,7 @@ func runBench(ctx *cli.Context, b bench.Benchmark) error {
 	}
 	var ui ui
 	if !globalQuiet && !globalJSON {
+		registerUI(&ui)
 		go ui.Run()
 	}
 
@@ -289,6 +290,7 @@ func runBench(ctx *cli.Context, b bench.Benchmark) error {
 	}
 	monitor.InfoLn("Cleanup Done.")
 	ui.Wait()
+	registerUI(nil)
 	return nil
 }
 

cli/print.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"unicode"
 
 	"github.com/cheggaaa/pb"
@@ -46,6 +47,13 @@ type errorMessage struct {
 
 var printMu sync.Mutex
 
+// Global UI instance for cleanup on fatal errors
+var globalUI atomic.Pointer[ui]
+
+func registerUI(u *ui) {
+	globalUI.Store(u)
+}
+
 func printError(data ...any) {
 	printMu.Lock()
 	defer printMu.Unlock()
@@ -63,6 +71,10 @@ func fatalIf(err *probe.Error, msg string, data ...any) {
 	if err == nil {
 		return
 	}
+	// Clean up UI if one is registered
+	if u := globalUI.Load(); u != nil {
+		u.Wait()
+	}
 	fatal(err, msg, data...)
 }