fix(ui): hold stream waits in tests and wait for slow work until it answers

AGENTS.md

@@ -33,7 +33,8 @@ the opposite caused a bug that took hours to find.
 3. **Polling lives in `tea.Cmd`.** A poll returns a message. A goroutine never
    writes to the model, never draws, and never blocks sending on a channel
    the UI also reads. A timer starts only through `startTimer`
-   (`internal/ui/timer.go`).
+   (`internal/ui/timer.go`). A command waits on a stream only through `waitOn`
+   (`internal/ui/stream.go`).
 4. **Events say when, reads say what.** The event stream marks a screen
    stale, and the screen is read again the usual way. A cluster that does not
    stream is polled as before. Event payloads are read in two places only:

internal/ui/feed.go

@@ -200,19 +200,15 @@ func nextEvent(feed *nomad.Feed) tea.Cmd {
 		return nil
 	}
 
-	return func() tea.Msg {
-		select {
-		case event, ok := <-feed.C:
-			if !ok {
-				return feedEndedMsg{feed: feed, err: reasonOf(feed)}
-			}
-
-			return feedEventsMsg{feed: feed, events: withWaiting(feed, event)}
-
-		case err := <-feed.Err:
-			return feedEndedMsg{feed: feed, err: err}
+	return waitOn(feed.C, feed.Err, func(event nomad.Event, ok bool) tea.Msg {
+		if !ok {
+			return feedEndedMsg{feed: feed, err: reasonOf(feed)}
 		}
-	}
+
+		return feedEventsMsg{feed: feed, events: withWaiting(feed, event)}
+	}, func(err error) tea.Msg {
+		return feedEndedMsg{feed: feed, err: err}
+	})
 }
 
 // withWaiting is the first event and every one waiting behind it, oldest

internal/ui/logs.go

@@ -256,23 +256,19 @@ func waitForStream(stream *nomad.LogStream) tea.Cmd {
 		return nil
 	}
 
-	return func() tea.Msg {
-		select {
-		case line, ok := <-stream.Lines:
-			if !ok {
-				return endOf(stream)
-			}
-
-			return logLineMsg{stream: stream, text: line, start: stream.Start(), old: stream.Old(line)}
-
-		case err := <-stream.Err:
-			if err != nil {
-				return errMsg{err: err}
-			}
+	return waitOn(stream.Lines, stream.Err, func(line string, ok bool) tea.Msg {
+		if !ok {
+			return endOf(stream)
+		}
 
-			return logEndMsg{stream: stream}
+		return logLineMsg{stream: stream, text: line, start: stream.Start(), old: stream.Old(line)}
+	}, func(err error) tea.Msg {
+		if err != nil {
+			return errMsg{err: err}
 		}
-	}
+
+		return logEndMsg{stream: stream}
+	})
 }
 
 // endOf is the end of a stream whose lines closed: its error, when it sent

internal/ui/stream.go

@@ -0,0 +1,39 @@
+package ui
+
+import (
+	"sync/atomic"
+
+	tea "charm.land/bubbletea/v2"
+)
+
+// heldStream is what a held stream wait answers when nothing has come: it
+// would wait.
+type heldStream struct{}
+
+// streamsHeld is set by tests that play out commands, to tell a stream wait
+// from slow work.
+var streamsHeld atomic.Bool
+
+// waitOn is the one place a command waits on a stream: for the next thing
+// on c, or for the error the stream ends with. A held wait does not wait.
+func waitOn[T any](c <-chan T, errs <-chan error, next func(T, bool) tea.Msg, ended func(error) tea.Msg) tea.Cmd {
+	return func() tea.Msg {
+		if streamsHeld.Load() {
+			select {
+			case v, ok := <-c:
+				return next(v, ok)
+			case err := <-errs:
+				return ended(err)
+			default:
+				return heldStream{}
+			}
+		}
+
+		select {
+		case v, ok := <-c:
+			return next(v, ok)
+		case err := <-errs:
+			return ended(err)
+		}
+	}
+}

internal/ui/stream_test.go

@@ -0,0 +1,119 @@
+package ui
+
+import (
+	"go/ast"
+	"go/token"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ingvarch/urga/internal/nomad"
+)
+
+// waitsOnChannels are the functions that may wait on a channel: the one
+// place a command waits on a stream, and the shell's watch of the terminal
+// size, which runs in the shell's own goroutine and not in a command.
+var waitsOnChannels = map[string]bool{"waitOn": true, "watchSize": true, "sendSize": true}
+
+func TestStreams_AreWaitedOnInOnePlace(t *testing.T) {
+	fset, files := sourceFiles(t)
+
+	for _, file := range files {
+		for _, decl := range file.Decls {
+			fn, ok := decl.(*ast.FuncDecl)
+			if !ok || fn.Body == nil || waitsOnChannels[fn.Name.Name] {
+				continue
+			}
+
+			for _, at := range channelWaits(fn.Body) {
+				t.Errorf("%s waits on a channel outside waitOn: %s", fn.Name.Name, fset.Position(at))
+			}
+		}
+	}
+}
+
+// channelWaits are the receives in body that can block: every receive but
+// the cases of a select that has a default. A range over a channel waits too,
+// but finding one needs the types; playOut names such a command when it
+// does not answer.
+func channelWaits(body *ast.BlockStmt) []token.Pos {
+	nonBlocking := map[ast.Node]bool{}
+
+	ast.Inspect(body, func(n ast.Node) bool {
+		sel, ok := n.(*ast.SelectStmt)
+		if !ok || !hasDefault(sel) {
+			return true
+		}
+
+		for _, stmt := range sel.Body.List {
+			if comm := stmt.(*ast.CommClause).Comm; comm != nil {
+				ast.Inspect(comm, func(n ast.Node) bool {
+					if recv, ok := n.(*ast.UnaryExpr); ok && recv.Op == token.ARROW {
+						nonBlocking[recv] = true
+					}
+
+					return true
+				})
+			}
+		}
+
+		return true
+	})
+
+	var waits []token.Pos
+
+	ast.Inspect(body, func(n ast.Node) bool {
+		if recv, ok := n.(*ast.UnaryExpr); ok && recv.Op == token.ARROW && !nonBlocking[recv] {
+			waits = append(waits, recv.Pos())
+		}
+
+		return true
+	})
+
+	return waits
+}
+
+func hasDefault(sel *ast.SelectStmt) bool {
+	for _, stmt := range sel.Body.List {
+		if stmt.(*ast.CommClause).Comm == nil {
+			return true
+		}
+	}
+
+	return false
+}
+
+// holdStreams makes every stream wait answer at once, for the length of the
+// test.
+func holdStreams(t *testing.T) {
+	t.Helper()
+
+	held := streamsHeld.Swap(true)
+
+	t.Cleanup(func() { streamsHeld.Store(held) })
+}
+
+func TestStreams_AHeldWaitAnswersWithWhatCameOrAtOnce(t *testing.T) {
+	r := require.New(t)
+
+	holdStreams(t)
+
+	lines := make(chan string, 1)
+	lines <- "the task says something\n"
+	stream := &nomad.LogStream{Lines: lines, Err: make(chan error)}
+
+	msg, ok := within(waitForStream(stream), time.Second).(logLineMsg)
+	r.True(ok)
+	r.Equal("the task says something\n", msg.text)
+
+	// Nothing more came: the wait answers that it would wait.
+	r.Equal(heldStream{}, within(waitForStream(stream), time.Second))
+}
+
+func TestPlayOut_LeavesAStreamThatSentNothing(t *testing.T) {
+	m := playOut(loadedModel(t), waitForStream(&nomad.LogStream{Lines: make(chan string)}))
+
+	require.Contains(t, plain(m.render()), "web")
+	require.False(t, streamsHeld.Load())
+}

internal/ui/watch.go

@@ -107,22 +107,18 @@ func (w watchState) waitForChange() tea.Cmd {
 
 	id := w.id
 
-	return func() tea.Msg {
-		select {
-		case _, ok := <-changes.C:
-			if !ok {
-				// A stream that dropped sends its error before it
-				// closes, and both are then ready at once: reading the
-				// close first must not lose the error.
-				return watchEndedMsg{id: id, err: reasonOf(changes)}
-			}
-
-			return changeMsg{id: id}
-
-		case err := <-changes.Err:
-			return watchEndedMsg{id: id, err: err}
+	return waitOn(changes.C, changes.Err, func(_ nomad.Change, ok bool) tea.Msg {
+		if !ok {
+			// A stream that dropped sends its error before it
+			// closes, and both are then ready at once: reading the
+			// close first must not lose the error.
+			return watchEndedMsg{id: id, err: reasonOf(changes)}
 		}
-	}
+
+		return changeMsg{id: id}
+	}, func(err error) tea.Msg {
+		return watchEndedMsg{id: id, err: err}
+	})
 }
 
 // reasonOf is the error a stream ended with, if it sent one.