Try to find misuses of textui.Progress

 - Add a runtime-check to Progress to notice if we deadlocked or
   forgot to call .Done().

 - Add a runtime-check to Progress.Done() to panic if .Set() was never
   called (instead of the old behavior of deadlocking).

 - grep: Use `defer` when possible, to help remember to call .Done().

 - grep: Always either call .Set() right away, or right before calling
   .Done().

4 files changed, 108 insertions, 58 deletions

diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index 8e26c08..860a49c 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -261,68 +261,75 @@ func (g Graph) InsertNode(node *btrfstree.Node) {
 }
 
 func (g Graph) FinalCheck(ctx context.Context, fs btrfstree.NodeSource) error {
-	var stats textui.Portion[int]
+	{
+		dlog.Info(ctx, "Checking keypointers for dead-ends...")
 
-	dlog.Info(ctx, "Checking keypointers for dead-ends...")
-	progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	stats.D = len(g.EdgesTo)
-	progressWriter.Set(stats)
-	for laddr := range g.EdgesTo {
-		if !maps.HasKey(g.Nodes, laddr) {
-			node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
-				LAddr: containers.OptionalValue(laddr),
-			})
-			fs.ReleaseNode(node)
-			if err == nil {
-				progressWriter.Done()
-				return fmt.Errorf("node@%v exists but was not in node scan results", laddr)
-			}
-			g.BadNodes[laddr] = err
-		}
-		stats.N++
+		var stats textui.Portion[int]
+		stats.D = len(g.EdgesTo)
+		progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 		progressWriter.Set(stats)
-	}
-	progressWriter.Done()
-	dlog.Info(ctx, "... done checking keypointers")
 
-	dlog.Info(ctx, "Checking for btree loops...")
-	stats.D = len(g.Nodes)
-	stats.N = 0
-	progressWriter = textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	progressWriter.Set(stats)
-	visited := make(containers.Set[btrfsvol.LogicalAddr], len(g.Nodes))
-	numLoops := 0
-	var checkNode func(node btrfsvol.LogicalAddr, stack []btrfsvol.LogicalAddr)
-	checkNode = func(node btrfsvol.LogicalAddr, stack []btrfsvol.LogicalAddr) {
-		defer func() {
-			stats.N = len(visited)
+		for laddr := range g.EdgesTo {
+			if !maps.HasKey(g.Nodes, laddr) {
+				node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
+					LAddr: containers.OptionalValue(laddr),
+				})
+				fs.ReleaseNode(node)
+				if err == nil {
+					progressWriter.Done()
+					return fmt.Errorf("node@%v exists but was not in node scan results", laddr)
+				}
+				g.BadNodes[laddr] = err
+			}
+			stats.N++
 			progressWriter.Set(stats)
-		}()
-		if visited.Has(node) {
-			return
 		}
-		if slices.Contains(node, stack) {
-			numLoops++
-			dlog.Error(ctx, "loop:")
-			for _, line := range g.renderLoop(append(stack, node)) {
-				dlog.Errorf(ctx, "    %s", line)
+		progressWriter.Done()
+		dlog.Info(ctx, "... done checking keypointers")
+	}
+
+	{
+		dlog.Info(ctx, "Checking for btree loops...")
+
+		var stats textui.Portion[int]
+		stats.D = len(g.Nodes)
+		progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+		progressWriter.Set(stats)
+
+		visited := make(containers.Set[btrfsvol.LogicalAddr], len(g.Nodes))
+		numLoops := 0
+		var checkNode func(node btrfsvol.LogicalAddr, stack []btrfsvol.LogicalAddr)
+		checkNode = func(node btrfsvol.LogicalAddr, stack []btrfsvol.LogicalAddr) {
+			defer func() {
+				stats.N = len(visited)
+				progressWriter.Set(stats)
+			}()
+			if visited.Has(node) {
+				return
+			}
+			if slices.Contains(node, stack) {
+				numLoops++
+				dlog.Error(ctx, "loop:")
+				for _, line := range g.renderLoop(append(stack, node)) {
+					dlog.Errorf(ctx, "    %s", line)
+				}
+				return
 			}
-			return
+			stack = append(stack, node)
+			for _, kp := range g.EdgesTo[node] {
+				checkNode(kp.FromNode, stack)
+			}
+			visited.Insert(node)
 		}
-		stack = append(stack, node)
-		for _, kp := range g.EdgesTo[node] {
-			checkNode(kp.FromNode, stack)
+		for _, node := range maps.SortedKeys(g.Nodes) {
+			checkNode(node, nil)
 		}
-		visited.Insert(node)
-	}
-	for _, node := range maps.SortedKeys(g.Nodes) {
-		checkNode(node, nil)
-	}
-	progressWriter.Done()
-	if numLoops > 0 {
-		return fmt.Errorf("%d btree loops", numLoops)
+		progressWriter.Done()
+		if numLoops > 0 {
+			return fmt.Errorf("%d btree loops", numLoops)
+		}
+		dlog.Info(ctx, "... done checking for loops")
 	}
-	dlog.Info(ctx, "... done checking for loops")
 
 	return nil
 }
@@ -352,6 +359,7 @@ func ReadGraph(_ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAd
 	progressWriter.Set(stats)
 	for _, laddr := range nodeList {
 		if err := ctx.Err(); err != nil {
+			progressWriter.Done()
 			return Graph{}, err
 		}
 		node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
@@ -359,6 +367,7 @@ func ReadGraph(_ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAd
 		})
 		if err != nil {
 			fs.ReleaseNode(node)
+			progressWriter.Done()
 			return Graph{}, err
 		}
 		graph.InsertNode(node)
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index ffb2e5f..31a31be 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -259,11 +259,9 @@ func (tree *RebuiltTree) items(ctx context.Context, inc bool) containers.SortedM
 			progressWriter.Set(stats)
 		}
 	}
-	if stats.Leafs.N > 0 {
-		stats.Leafs.N = stats.Leafs.D
-		progressWriter.Set(stats)
-		progressWriter.Done()
-	}
+	stats.Leafs.N = stats.Leafs.D
+	progressWriter.Set(stats)
+	progressWriter.Done()
 
 	return index
 }
diff --git a/lib/streamio/runescanner.go b/lib/streamio/runescanner.go
index 203439f..1206522 100644
--- a/lib/streamio/runescanner.go
+++ b/lib/streamio/runescanner.go
@@ -103,6 +103,7 @@ func (rs *runeScanner) UnreadRune() error {
 
 // ReadRune implements io.Closer.
 func (rs *runeScanner) Close() error {
+	rs.progressWriter.Set(rs.progress)
 	rs.progressWriter.Done()
 	return rs.closer.Close()
 }
diff --git a/lib/textui/progress.go b/lib/textui/progress.go
index 48a3901..04c8212 100644
--- a/lib/textui/progress.go
+++ b/lib/textui/progress.go
@@ -18,6 +18,16 @@ type Stats interface {
 	fmt.Stringer
 }
 
+// Progress helps display to the user the ongoing progress of a long
+// task.
+//
+// There are few usage requirements to watch out for:
+//
+//   - .Set() must have been called at least once before you call
+//     .Done().  The easiest way to ensure this is to call .Set right
+//     after creating the progress, or right before calling .Done().  I
+//     advise against counting on a loop to have called .Set() at least
+//     once.
 type Progress[T Stats] struct {
 	ctx      context.Context //nolint:containedctx // captured for separate goroutine
 	lvl      dlog.LogLevel
@@ -29,6 +39,10 @@ type Progress[T Stats] struct {
 	cur     typedsync.Value[T]
 	oldStat T
 	oldLine string
+
+	// This isn't a functional part, but is useful for helping us
+	// to detect misuse.
+	last time.Time
 }
 
 func NewProgress[T Stats](ctx context.Context, lvl dlog.LogLevel, interval time.Duration) *Progress[T] {
@@ -44,18 +58,43 @@ func NewProgress[T Stats](ctx context.Context, lvl dlog.LogLevel, interval time.
 	return ret
 }
 
+// Set update the Progress.  Rate-limiting prevents this from being
+// expensive, or from spamming the user; it is reasonably safe to call
+// .Set in a tight inner loop.
+//
+// It is safe to call Set concurrently.
 func (p *Progress[T]) Set(val T) {
 	if _, hadOld := p.cur.Swap(val); !hadOld {
 		go p.run()
 	}
 }
 
+// Done closes the Progress; it flushes out one last status update (if
+// nescessary), and releases resources associated with the Progress.
+//
+// It is safe to call Done multiple times, or concurrently.
+//
+// It will panic if Done is called without having called Set at least
+// once.
 func (p *Progress[T]) Done() {
 	p.cancel()
+	if _, started := p.cur.Load(); !started {
+		panic("textui.Progress: .Done called without ever calling .Set")
+	}
 	<-p.done
 }
 
 func (p *Progress[T]) flush(force bool) {
+	// Check how long it's been since we last printed something.
+	// If this grows too big, it probably means that either the
+	// program deadlocked or that we forgot to call .Done().
+	now := time.Now()
+	if !p.last.IsZero() && now.Sub(p.last) > Tunable(2*time.Minute) {
+		dlog.Error(p.ctx, "stale Progress")
+		panic("stale Progress")
+	}
+
+	// Load the data to print.
 	cur, ok := p.cur.Load()
 	if !ok {
 		panic("should not happen")
@@ -65,13 +104,16 @@ func (p *Progress[T]) flush(force bool) {
 	}
 	defer func() { p.oldStat = cur }()
 
+	// Format the data as text.
 	line := cur.String()
 	if !force && line == p.oldLine {
 		return
 	}
 	defer func() { p.oldLine = line }()
 
+	// Print.
 	dlog.Log(p.ctx, p.lvl, line)
+	p.last = now
 }
 
 func (p *Progress[T]) run() {