Merge branch 'lukeshu/fixes'

10 files changed, 145 insertions, 82 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
index 427070a..e6345ae 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
@@ -190,12 +190,12 @@ func (o *rebuilder) processAddedItemQueue(ctx context.Context) error {
 	var progress settleItemStats
 	progress.D = len(queue)
 	progressWriter := textui.NewProgress[settleItemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.substep.progress", &progress)
+	progressWriter.Set(progress)
+	defer progressWriter.Done()
 
-	for i, key := range queue {
-		progress.N = i
-		progressWriter.Set(progress)
+	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.substep.progress", &progress)
 
+	for _, key := range queue {
 		ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.settle.item", key)
 		tree := o.rebuilt.RebuiltTree(ctx, key.TreeID)
 		incPtr, ok := tree.RebuiltAcquireItems(ctx).Load(key.Key)
@@ -213,15 +213,14 @@ func (o *rebuilder) processAddedItemQueue(ctx context.Context) error {
 			o.wantAugment(ctx, wantKey, tree.RebuiltLeafToRoots(ctx, excPtr.Node))
 			progress.NumAugments = o.numAugments
 			progress.NumAugmentTrees = len(o.augmentQueue)
-			progressWriter.Set(progress)
 		} else if !btrfscheck.HandleItemWouldBeNoOp(key.ItemType) {
 			o.settledItemQueue.Insert(key)
 		}
+
+		progress.N++
+		progressWriter.Set(progress)
 	}
 
-	progress.N = len(queue)
-	progressWriter.Set(progress)
-	progressWriter.Done()
 	return nil
 }
 
@@ -251,6 +250,9 @@ func (o *rebuilder) processSettledItemQueue(ctx context.Context) error {
 	var progress processItemStats
 	progress.D = len(queue)
 	progressWriter := textui.NewProgress[processItemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progressWriter.Set(progress)
+	defer progressWriter.Done()
+
 	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.substep.progress", &progress)
 
 	type keyAndBody struct {
@@ -278,7 +280,6 @@ func (o *rebuilder) processSettledItemQueue(ctx context.Context) error {
 		return nil
 	})
 	grp.Go("cpu", func(ctx context.Context) error {
-		defer progressWriter.Done()
 		o.curKey.Key.OK = true
 		for item := range itemChan {
 			ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.process.item", item.keyAndTree)
@@ -323,24 +324,23 @@ func (o *rebuilder) processAugmentQueue(ctx context.Context) error {
 	runtime.GC()
 
 	progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progressWriter.Set(progress)
+	defer progressWriter.Done()
+
 	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.substep.progress", &progress)
 	for _, treeID := range maps.SortedKeys(resolvedAugments) {
 		ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.augment.tree", treeID)
 		for _, nodeAddr := range maps.SortedKeys(resolvedAugments[treeID]) {
 			if err := ctx.Err(); err != nil {
-				progressWriter.Set(progress)
-				progressWriter.Done()
 				return err
 			}
-			progressWriter.Set(progress)
 			// This will call o.AddedItem as nescessary, which
 			// inserts to o.addedItemQueue.
 			o.rebuilt.RebuiltTree(ctx, treeID).RebuiltAddRoot(ctx, nodeAddr)
 			progress.N++
+			progressWriter.Set(progress)
 		}
 	}
-	progressWriter.Set(progress)
-	progressWriter.Done()
 
 	return nil
 }
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
index f266dab..b1469ff 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
@@ -73,6 +73,7 @@ func ScanDevices(_ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.Logical
 	progressWriter.Set(stats)
 	for _, laddr := range nodeList {
 		if err := ctx.Err(); err != nil {
+			progressWriter.Done()
 			return ScanDevicesResult{}, err
 		}
 		node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
@@ -80,6 +81,7 @@ func ScanDevices(_ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.Logical
 		})
 		if err != nil {
 			fs.ReleaseNode(node)
+			progressWriter.Done()
 			return ScanDevicesResult{}, err
 		}
 		ret.insertNode(node)
diff --git a/lib/btrfs/btrfsitem/item_extent.go b/lib/btrfs/btrfsitem/item_extent.go
index edfef7c..cf9b739 100644
--- a/lib/btrfs/btrfsitem/item_extent.go
+++ b/lib/btrfs/btrfsitem/item_extent.go
@@ -54,7 +54,9 @@ func (o Extent) Clone() Extent {
 	ret.Refs = extentInlineRefPool.Get(len(o.Refs))
 	copy(ret.Refs, o.Refs)
 	for i := range ret.Refs {
-		ret.Refs[i].Body = o.Refs[i].Body.CloneItem()
+		if o.Refs[i].Body != nil {
+			ret.Refs[i].Body = o.Refs[i].Body.CloneItem()
+		}
 	}
 	return ret
 }
diff --git a/lib/btrfs/btrfsitem/item_metadata.go b/lib/btrfs/btrfsitem/item_metadata.go
index 10ae994..90f7622 100644
--- a/lib/btrfs/btrfsitem/item_metadata.go
+++ b/lib/btrfs/btrfsitem/item_metadata.go
@@ -46,7 +46,9 @@ func (o Metadata) Clone() Metadata {
 	ret.Refs = extentInlineRefPool.Get(len(o.Refs))
 	copy(ret.Refs, o.Refs)
 	for i := range ret.Refs {
-		ret.Refs[i].Body = o.Refs[i].Body.CloneItem()
+		if o.Refs[i].Body != nil {
+			ret.Refs[i].Body = o.Refs[i].Body.CloneItem()
+		}
 	}
 	return ret
 }
diff --git a/lib/btrfscheck/graph.go b/lib/btrfscheck/graph.go
index 806f609..bd6d1fe 100644
--- a/lib/btrfscheck/graph.go
+++ b/lib/btrfscheck/graph.go
@@ -258,7 +258,7 @@ func HandleItem(ctx context.Context, o GraphCallbacks, treeID btrfsprim.ObjID, i
 			case nil:
 				// nothing
 			case *btrfsitem.ExtentDataRef:
-				o.WantOff(ctx, "referencing INode",
+				o.WantOff(ctx, "referencing Inode",
 					refBody.Root,
 					refBody.ObjectID,
 					btrfsitem.INODE_ITEM_KEY,
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/maps/maputil.go b/lib/maps/maputil.go
index 63e52a0..3b6691a 100644
--- a/lib/maps/maputil.go
+++ b/lib/maps/maputil.go
@@ -31,13 +31,20 @@ func HasKey[K comparable, V any](m map[K]V, k K) bool {
 	return has
 }
 
-func HaveAnyKeysInCommon[K comparable, V1, V2 any](small map[K]V1, big map[K]V2) bool {
-	if len(big) < len(small) {
-		return HaveAnyKeysInCommon(big, small)
-	}
-	for v := range small {
-		if _, ok := big[v]; ok {
-			return true
+func HaveAnyKeysInCommon[K comparable, V1, V2 any](a map[K]V1, b map[K]V2) bool {
+	if len(a) < len(b) {
+		small, big := a, b
+		for v := range small {
+			if _, ok := big[v]; ok {
+				return true
+			}
+		}
+	} else {
+		small, big := b, a
+		for v := range small {
+			if _, ok := big[v]; ok {
+				return true
+			}
 		}
 	}
 	return false
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() {