From 3609d87e703cf7f424fa70e72515a8472b843eaf Mon Sep 17 00:00:00 2001
From: Christian Muehlhaeuser <muesli@gmail.com>
Date: Sat, 15 Oct 2022 04:04:06 +0200
Subject: [PATCH] fix: don't block in Send after shutdown

Send should block before a tea.Program has been started, but result
in a no-op when it has already been terminated.

Fixed godocs.
---
 tea.go      | 26 ++++++++++++--------------
 tea_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/tea.go b/tea.go
index 1519ade..d6ba73c 100644
--- a/tea.go
+++ b/tea.go
@@ -142,6 +142,9 @@ func NewProgram(model Model, opts ...ProgramOption) *Program {
 		msgs:         make(chan Msg),
 	}
 
+	// Initialize context and teardown channel.
+	p.ctx, p.cancel = context.WithCancel(context.Background())
+
 	// Apply all options to the program.
 	for _, opt := range opts {
 		opt(p)
@@ -246,12 +249,7 @@ func (p *Program) handleCommands(cmds chan Cmd) chan struct{} {
 				// (e.g. tick commands that sleep for half a second). It's not
 				// possible to cancel them so we'll have to leak the goroutine
 				// until Cmd returns.
-				go func() {
-					select {
-					case p.msgs <- cmd():
-					case <-p.ctx.Done():
-					}
-				}()
+				go p.Send(cmd())
 			}
 		}
 	}()
@@ -315,10 +313,7 @@ func (p *Program) eventLoop(model Model, cmds chan Cmd) (Model, error) {
 				go func() {
 					// Execute commands one at a time, in order.
 					for _, cmd := range msg {
-						select {
-						case p.msgs <- cmd():
-						case <-p.ctx.Done():
-						}
+						p.Send(cmd())
 					}
 				}()
 			}
@@ -344,7 +339,6 @@ func (p *Program) Run() (Model, error) {
 	cmds := make(chan Cmd)
 	p.errs = make(chan error)
 
-	p.ctx, p.cancel = context.WithCancel(context.Background())
 	defer p.cancel()
 
 	switch {
@@ -504,10 +498,14 @@ func (p *Program) Start() error {
 // messages to be injected from outside the program for interoperability
 // purposes.
 //
-// If the program is not running this will be a no-op, so it's safe to
-// send messages if the program is unstarted, or has exited.
+// If the program hasn't started yet this will be a blocking operation.
+// If the program has already been terminated this will be a no-op, so it's safe
+// to send messages after the program has exited.
 func (p *Program) Send(msg Msg) {
-	p.msgs <- msg
+	select {
+	case <-p.ctx.Done():
+	case p.msgs <- msg:
+	}
 }
 
 // Quit is a convenience function for quitting Bubble Tea programs. Use it
diff --git a/tea_test.go b/tea_test.go
index c61c518..4b087bd 100644
--- a/tea_test.go
+++ b/tea_test.go
@@ -149,3 +149,29 @@ func TestTeaSequenceMsg(t *testing.T) {
 		t.Fatalf("counter should be 2, got %d", m.counter.Load())
 	}
 }
+
+func TestTeaSend(t *testing.T) {
+	var buf bytes.Buffer
+	var in bytes.Buffer
+
+	m := &testModel{}
+	p := NewProgram(m, WithInput(&in), WithOutput(&buf))
+
+	// sending before the program is started is a blocking operation
+	go p.Send(Quit())
+
+	if _, err := p.Run(); err != nil {
+		t.Fatal(err)
+	}
+
+	// sending a message after program has quit is a no-op
+	p.Send(Quit())
+}
+
+func TestTeaNoRun(t *testing.T) {
+	var buf bytes.Buffer
+	var in bytes.Buffer
+
+	m := &testModel{}
+	NewProgram(m, WithInput(&in), WithOutput(&buf))
+}