Repository: charmbracelet/bubbletea
Pull request #1140 (2024-09-12 16:55:52)

fix: use safeWriter to guard writing to output
This implements a wrapper around p.output that locks the output on every write. We also need to redefine `channelHandlers` to be thread-safe since `p.Kill()` can be called from a different thread while the program is running.


Patch:
diff --git a/examples/go.mod b/examples/go.mod
index 5c1dccd..19d2cab 100644
--- a/examples/go.mod
+++ b/examples/go.mod
@@ -22,7 +22,7 @@ require (
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/aymanbagabas/go-udiff v0.2.0 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
-	github.com/charmbracelet/x/ansi v0.2.3 // indirect
+	github.com/charmbracelet/x/ansi v0.3.0 // indirect
 	github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b // indirect
 	github.com/charmbracelet/x/term v0.2.0 // indirect
 	github.com/charmbracelet/x/windows v0.2.0 // indirect
diff --git a/examples/go.sum b/examples/go.sum
index 6c9a51a..d2f0383 100644
--- a/examples/go.sum
+++ b/examples/go.sum
@@ -14,8 +14,6 @@ github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWp
 github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
-github.com/charmbracelet/bubbles v0.20.0 h1:jSZu6qD8cRQ6k9OMfR1WlM+ruM8fkPWkHvQWD9LIutE=
-github.com/charmbracelet/bubbles v0.20.0/go.mod h1:39slydyswPy+uVOHZ5x/GjwVAFkCsV8IIVy+4MhzwwU=
 github.com/charmbracelet/bubbles v0.20.1-0.20240910180436-8972b56c9dde h1:/u+F3xZsfJYhEQ3MjKvqiXb2Y6LV8/4YHAMICxjGVhY=
 github.com/charmbracelet/bubbles v0.20.1-0.20240910180436-8972b56c9dde/go.mod h1:M5Z4yZGaxq8TZoEY2rs529UxFTeZLnAhmAzqd2wd48M=
 github.com/charmbracelet/glamour v0.8.0 h1:tPrjL3aRcQbn++7t18wOpgLyl8wrOHUEDS7IZ68QtZs=
@@ -24,12 +22,10 @@ github.com/charmbracelet/harmonica v0.2.0 h1:8NxJWRWg/bzKqqEaaeFNipOu77YR5t8aSwG
 github.com/charmbracelet/harmonica v0.2.0/go.mod h1:KSri/1RMQOZLbw7AHqgcBycp8pgJnQMYYT8QZRqZ1Ao=
 github.com/charmbracelet/lipgloss v0.13.0 h1:4X3PPeoWEDCMvzDvGmTajSyYPcZM4+y8sCA/SsA3cjw=
 github.com/charmbracelet/lipgloss v0.13.0/go.mod h1:nw4zy0SBX/F/eAO1cWdcvy6qnkDUxr8Lw7dvFrAIbbY=
-github.com/charmbracelet/x/ansi v0.2.3 h1:VfFN0NUpcjBRd4DnKfRaIRo53KRgey/nhOoEqosGDEY=
-github.com/charmbracelet/x/ansi v0.2.3/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw=
+github.com/charmbracelet/x/ansi v0.3.0 h1:CCsscv7vKC/DNYUYFQNNIOWzrpTUbLXL3d4fdFIQ0WE=
+github.com/charmbracelet/x/ansi v0.3.0/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw=
 github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b h1:MnAMdlwSltxJyULnrYbkZpp4k58Co7Tah3ciKhSNo0Q=
 github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
-github.com/charmbracelet/x/exp/teatest v0.0.0-20240521184646-23081fb03b28 h1:sOWKNRjt8uOEVgPiJVIJCse1+mUDM2F/vYY6W0Go640=
-github.com/charmbracelet/x/exp/teatest v0.0.0-20240521184646-23081fb03b28/go.mod h1:l1w+LTJZCCozeGzMEWGxRw6Mo2DfcZUvupz8HGubdes=
 github.com/charmbracelet/x/exp/teatest v0.0.0-20240912153648-d6041061ead9 h1:2q41wBBupsw4GATbF/C3NZ3JtccZ/iYu7cL3doHOnuo=
 github.com/charmbracelet/x/exp/teatest v0.0.0-20240912153648-d6041061ead9/go.mod h1:WPIoTqLbtPGmiv83mabqvIYOrwa1vUZMdp3urF4gQrk=
 github.com/charmbracelet/x/term v0.2.0 h1:cNB9Ot9q8I711MyZ7myUR5HFWL/lc3OpU8jZ4hwm0x0=
diff --git a/exec.go b/exec.go
index 7a14d2a..ef353eb 100644
--- a/exec.go
+++ b/exec.go
@@ -109,7 +109,7 @@ func (p *Program) exec(c ExecCommand, fn ExecCallback) {
 	}
 
 	c.SetStdin(p.input)
-	c.SetStdout(p.output)
+	c.SetStdout(p.output.Writer())
 	c.SetStderr(os.Stderr)
 
 	// Execute system command.
diff --git a/options.go b/options.go
index ea43fd7..af00f32 100644
--- a/options.go
+++ b/options.go
@@ -29,7 +29,7 @@ func WithContext(ctx context.Context) ProgramOption {
 // won't need to use this.
 func WithOutput(output io.Writer) ProgramOption {
 	return func(p *Program) {
-		p.output = output
+		p.output = newSafeWriter(output)
 	}
 }
 
diff --git a/standard_renderer.go b/standard_renderer.go
index ae2d439..e59023d 100644
--- a/standard_renderer.go
+++ b/standard_renderer.go
@@ -64,19 +64,18 @@ func (r *standardRenderer) setOutput(out io.Writer) {
 
 // close closes the renderer and flushes any remaining data.
 func (r *standardRenderer) close() (err error) {
-	r.mtx.Lock()
-	defer r.mtx.Unlock()
-
-	r.execute(ansi.EraseEntireLine)
 	// Move the cursor back to the beginning of the line
-	r.execute("\r")
+	// NOTE: execute locks the mutex
+	r.execute(ansi.EraseEntireLine + "\r")
 
 	return
 }
 
 // execute writes the given sequence to the output.
 func (r *standardRenderer) execute(seq string) {
+	r.mtx.Lock()
 	_, _ = io.WriteString(r.out, seq)
+	r.mtx.Unlock()
 }
 
 // flush renders the buffer.
@@ -240,11 +239,7 @@ func (r *standardRenderer) reset() {
 }
 
 func (r *standardRenderer) clearScreen() {
-	r.mtx.Lock()
-	defer r.mtx.Unlock()
-
-	r.execute(ansi.EraseEntireDisplay)
-	r.execute(ansi.MoveCursorOrigin)
+	r.execute(ansi.EraseEntireDisplay + ansi.MoveCursorOrigin)
 
 	r.repaint()
 }
diff --git a/sync.go b/sync.go
new file mode 100644
index 0000000..e8cd21f
--- /dev/null
+++ b/sync.go
@@ -0,0 +1,31 @@
+package tea
+
+import (
+	"io"
+	"sync"
+)
+
+// safeWriter is a thread-safe writer.
+type safeWriter struct {
+	w  io.Writer
+	mu sync.Mutex
+}
+
+var _ io.Writer = &safeWriter{}
+
+// newSafeWriter returns a new safeWriter.
+func newSafeWriter(w io.Writer) *safeWriter {
+	return &safeWriter{w: w}
+}
+
+// Writer returns the underlying writer.
+func (w *safeWriter) Writer() io.Writer {
+	return w.w
+}
+
+// Write writes to the underlying writer.
+func (w *safeWriter) Write(p []byte) (n int, err error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.w.Write(p)
+}
diff --git a/tea.go b/tea.go
index 04c1e9b..89ff08e 100644
--- a/tea.go
+++ b/tea.go
@@ -109,18 +109,27 @@ const (
 // channelHandlers manages the series of channels returned by various processes.
 // It allows us to wait for those processes to terminate before exiting the
 // program.
-type channelHandlers []chan struct{}
+type channelHandlers struct {
+	handlers []chan struct{}
+	mu       sync.RWMutex
+}
 
 // Adds a channel to the list of handlers. We wait for all handlers to terminate
 // gracefully on shutdown.
 func (h *channelHandlers) add(ch chan struct{}) {
-	*h = append(*h, ch)
+	h.mu.Lock()
+	h.handlers = append(h.handlers, ch)
+	h.mu.Unlock()
 }
 
 // shutdown waits for all handlers to terminate.
-func (h channelHandlers) shutdown() {
+func (h *channelHandlers) shutdown() {
 	var wg sync.WaitGroup
-	for _, ch := range h {
+
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	for _, ch := range h.handlers {
 		wg.Add(1)
 		go func(ch chan struct{}) {
 			<-ch
@@ -151,12 +160,14 @@ type Program struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	msgs     chan Msg
-	errs     chan error
-	finished chan struct{}
+	msgs         chan Msg
+	errs         chan error
+	finished     chan struct{}
+	shutdownOnce sync.Once
 
 	// where to send output, this will usually be os.Stdout.
-	output io.Writer
+	output *safeWriter
+
 	// ttyOutput is null if output is not a TTY.
 	ttyOutput           term.File
 	previousOutputState *term.State
@@ -248,7 +259,7 @@ func NewProgram(model Model, opts ...ProgramOption) *Program {
 
 	// if no output was set, set it to stdout
 	if p.output == nil {
-		p.output = os.Stdout
+		p.output = newSafeWriter(os.Stdout)
 	}
 
 	// if no environment was set, set it to os.Environ()
@@ -616,7 +627,8 @@ func (p *Program) Run() (Model, error) {
 	}
 
 	// If no renderer is set use the standard one.
-	output := p.output
+	var output io.Writer
+	output = p.output
 	if p.renderer == nil {
 		// TODO(v2): remove the ANSI compressor
 		if p.startupOptions.has(withANSICompressor) {
@@ -809,30 +821,32 @@ func (p *Program) execute(seq string) {
 // shutdown performs operations to free up resources and restore the terminal
 // to its original state.
 func (p *Program) shutdown(kill bool) {
-	p.cancel()
-
-	// Wait for all handlers to finish.
-	p.handlers.shutdown()
-
-	// Check if the cancel reader has been setup before waiting and closing.
-	if p.inputReader != nil {
-		// Wait for input loop to finish.
-		if p.inputReader.Cancel() {
-			if !kill {
-				p.waitForReadLoop()
+	p.shutdownOnce.Do(func() {
+		p.cancel()
+
+		// Wait for all handlers to finish.
+		p.handlers.shutdown()
+
+		// Check if the cancel reader has been setup before waiting and closing.
+		if p.inputReader != nil {
+			// Wait for input loop to finish.
+			if p.inputReader.Cancel() {
+				if !kill {
+					p.waitForReadLoop()
+				}
 			}
+			_ = p.inputReader.Close()
 		}
-		_ = p.inputReader.Close()
-	}
 
-	if p.renderer != nil {
-		p.stopRenderer(kill)
-	}
+		if p.renderer != nil {
+			p.stopRenderer(kill)
+		}
 
-	_ = p.restoreTerminalState()
-	if !kill {
-		p.finished <- struct{}{}
-	}
+		_ = p.restoreTerminalState()
+		if !kill {
+			p.finished <- struct{}{}
+		}
+	})
 }
 
 // recoverFromPanic recovers from a panic, prints the stack trace, and restores
@@ -995,7 +1009,7 @@ func (p *Program) stopRenderer(kill bool) {
 	p.renderer.close() //nolint:errcheck
 
 	if !kill && p.startupOptions.has(withANSICompressor) {
-		if w, ok := p.output.(io.WriteCloser); ok {
+		if w, ok := p.output.Writer().(io.WriteCloser); ok {
 			_ = w.Close()
 		}
 	}
diff --git a/tty_unix.go b/tty_unix.go
index 5cbb4fe..36c9b36 100644
--- a/tty_unix.go
+++ b/tty_unix.go
@@ -22,7 +22,7 @@ func (p *Program) initInput() (err error) {
 		}
 	}
 
-	if f, ok := p.output.(term.File); ok && term.IsTerminal(f.Fd()) {
+	if f, ok := p.output.Writer().(term.File); ok && term.IsTerminal(f.Fd()) {
 		p.ttyOutput = f
 	}
 
diff --git a/tty_windows.go b/tty_windows.go
index a3a2525..3f0c700 100644
--- a/tty_windows.go
+++ b/tty_windows.go
@@ -34,7 +34,7 @@ func (p *Program) initInput() (err error) {
 	}
 
 	// Save output screen buffer state and enable VT processing.
-	if f, ok := p.output.(term.File); ok && term.IsTerminal(f.Fd()) {
+	if f, ok := p.output.Writer().(term.File); ok && term.IsTerminal(f.Fd()) {
 		p.ttyOutput = f
 		p.previousOutputState, err = term.GetState(f.Fd())
 		if err != nil {


Test patch:
diff --git a/options_test.go b/options_test.go
index ce8d417..b013c7d 100644
--- a/options_test.go
+++ b/options_test.go
@@ -11,7 +11,7 @@ func TestOptions(t *testing.T) {
 	t.Run("output", func(t *testing.T) {
 		var b bytes.Buffer
 		p := NewProgram(nil, WithOutput(&b))
-		if f, ok := p.output.(*os.File); ok {
+		if f, ok := p.output.Writer().(*os.File); ok {
 			t.Errorf("expected output to custom, got %v", f.Fd())
 		}
 	})