rebuildnodes: Integrate loop-checking in to Graph.FinalCheck

4 files changed, 82 insertions, 56 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index 50c75a4..ab4b55f 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -345,6 +345,8 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd
 		return
 	}
 	if slices.Contains(node, stack) {
+		// This is a panic because graph.FinalCheck() should
+		// have already checked for loops.
 		panic("loop")
 	}
 	if !tree.isOwnerOK(graph.Nodes[node].Owner, graph.Nodes[node].Generation) {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
index 1180b0d..d3dd19a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
@@ -5,8 +5,12 @@
 package graph
 
 import (
+	"context"
 	"fmt"
 	"reflect"
+	"time"
+
+	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
@@ -14,6 +18,9 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 type Node struct {
@@ -131,7 +138,7 @@ func New(sb btrfstree.Superblock) *Graph {
 	return g
 }
 
-func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
+func (g Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
 	nodeData := Node{
 		Level:      nodeRef.Data.Head.Level,
 		Generation: nodeRef.Data.Head.Generation,
@@ -184,23 +191,77 @@ func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.N
 	}
 }
 
-func (g *Graph) FinalCheck(fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, progress func(n, d int)) error {
-	total := len(g.EdgesTo)
-	done := 0
-	progress(done, total)
+type progStats struct {
+	textui.Portion[int]
+}
+
+func (s progStats) String() string {
+	return "... " + s.Portion.String()
+}
+
+func (g Graph) FinalCheck(ctx context.Context, fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock) error {
+	var stats progStats
+
+	dlog.Info(ctx, "Checking keypointers for dead-ends...")
+	progressWriter := textui.NewProgress[progStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+	stats.D = len(g.EdgesTo)
+	progressWriter.Set(stats)
 	for laddr := range g.EdgesTo {
 		if _, ok := g.Nodes[laddr]; !ok {
 			_, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, sb, laddr, btrfstree.NodeExpectations{
 				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			})
 			if err == nil {
+				progressWriter.Done()
 				return fmt.Errorf("node@%v exists but was not in node scan results", laddr)
 			}
 			g.BadNodes[laddr] = err
 		}
-		done++
-		progress(done, total)
+		stats.N++
+		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[progStats](ctx, dlog.LogLevelInfo, 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
+		}
+		stack = append(stack, node)
+		for _, kp := range g.EdgesTo[node] {
+			checkNode(kp.FromNode, stack)
+		}
+		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)
+	}
+	dlog.Info(ctx, "... done checking for loops")
+
 	return nil
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go
index e76985c..0e51805 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go
@@ -6,47 +6,18 @@ package graph
 
 import (
 	"fmt"
-	"io"
 	"strings"
 
 	"github.com/datawire/dlib/derror"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
-	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
-func (g Graph) ShowLoops(out io.Writer) {
-	orphans := make(containers.Set[btrfsvol.LogicalAddr])
-	for node := range g.Nodes {
-		if len(g.EdgesTo[node]) == 0 {
-			orphans.Insert(node)
-		}
-	}
-
-	loopWalk(out, g, 0)
-	for _, orphan := range maps.SortedKeys(orphans) {
-		loopWalk(out, g, orphan)
-	}
-}
-
-func loopWalk(out io.Writer, scanData Graph, stack ...btrfsvol.LogicalAddr) {
-	for _, kp := range scanData.EdgesFrom[stack[len(stack)-1]] {
-		childStack := append(stack, kp.ToNode)
-		if slices.Contains(kp.ToNode, stack) {
-			loopRender(out, scanData, childStack...)
-		} else {
-			loopWalk(out, scanData, childStack...)
-		}
-	}
-}
-
-func nodeRender(scanData Graph, node btrfsvol.LogicalAddr) []string {
+func (g Graph) renderNode(node btrfsvol.LogicalAddr) []string {
 	if node == 0 {
 		return []string{"root"}
-	} else if nodeData, ok := scanData.Nodes[node]; ok {
+	} else if nodeData, ok := g.Nodes[node]; ok {
 		return []string{
 			fmt.Sprintf("{addr:      %v,", node),
 			fmt.Sprintf(" level:     %v,", nodeData.Level),
@@ -61,7 +32,7 @@ func nodeRender(scanData Graph, node btrfsvol.LogicalAddr) []string {
 				nodeData.MaxItem.ItemType,
 				nodeData.MaxItem.Offset),
 		}
-	} else if nodeErr, ok := scanData.BadNodes[node]; ok {
+	} else if nodeErr, ok := g.BadNodes[node]; ok {
 		return []string{
 			fmt.Sprintf("{addr:%v,", node),
 			fmt.Sprintf(" err:%q}", nodeErr.Error()),
@@ -71,7 +42,7 @@ func nodeRender(scanData Graph, node btrfsvol.LogicalAddr) []string {
 	}
 }
 
-func edgeRender(scanData Graph, kp Edge) []string {
+func (g Graph) renderEdge(kp Edge) []string {
 	a := fmt.Sprintf("[%d]={", kp.FromItem)
 	b := strings.Repeat(" ", len(a))
 	ret := []string{
@@ -85,8 +56,8 @@ func edgeRender(scanData Graph, kp Edge) []string {
 	}
 
 	var err error
-	if toNode, ok := scanData.Nodes[kp.ToNode]; !ok {
-		err = scanData.BadNodes[kp.ToNode]
+	if toNode, ok := g.Nodes[kp.ToNode]; !ok {
+		err = g.BadNodes[kp.ToNode]
 	} else {
 		err = checkNodeExpectations(kp, toNode)
 	}
@@ -100,7 +71,7 @@ func edgeRender(scanData Graph, kp Edge) []string {
 	return ret
 }
 
-func loopRender(out io.Writer, scanData Graph, stack ...btrfsvol.LogicalAddr) {
+func (g Graph) renderLoop(stack []btrfsvol.LogicalAddr) []string {
 	var lines []string
 	add := func(suffixes []string) {
 		curLen := 0
@@ -126,20 +97,17 @@ func loopRender(out io.Writer, scanData Graph, stack ...btrfsvol.LogicalAddr) {
 
 	for i, node := range stack {
 		if i > 0 {
-			for _, kp := range scanData.EdgesTo[node] {
+			for _, kp := range g.EdgesTo[node] {
 				if kp.FromNode == stack[i-1] {
-					add(edgeRender(scanData, *kp))
+					add(g.renderEdge(*kp))
 					break
 				}
 			}
 		}
-		add(nodeRender(scanData, node))
+		add(g.renderNode(node))
 	}
 
-	fmt.Fprintln(out, "loop:")
-	for _, line := range lines {
-		fmt.Fprintln(out, "    "+line)
-	}
+	return lines
 }
 
 func checkNodeExpectations(kp Edge, toNode Node) error {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 5c2d0fd..cdd5cba 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -72,14 +72,9 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	progressWriter.Done()
 	dlog.Info(ctx, "... done reading node data")
 
-	progressWriter = textui.NewProgress[scanStats](ctx, dlog.LogLevelInfo, 1*time.Second)
-	dlog.Infof(ctx, "Checking keypointers for dead-ends...")
-	if err := nodeGraph.FinalCheck(fs, *sb, progress); err != nil {
+	if err := nodeGraph.FinalCheck(ctx, fs, *sb); err != nil {
 		return graph.Graph{}, nil, err
 	}
-	progressWriter.Done()
-	dlog.Info(ctx, "... done checking keypointers")
-
 	keyIO.SetGraph(*nodeGraph)
 
 	return *nodeGraph, keyIO, nil