btrfsutil: Add a ReadGraph function, rebuildtrees: Clean up scan to match

3 files changed, 83 insertions, 18 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
index a6f9c7a..f266dab 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
@@ -41,31 +41,35 @@ type ScanDevicesResult struct {
 	Sizes map[btrfsutil.ItemPtr]SizeAndErr  // EXTENT_CSUM and EXTENT_DATA
 }
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr) (ScanDevicesResult, error) {
+func ScanDevices(_ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr) (ScanDevicesResult, error) {
+	// read-superblock /////////////////////////////////////////////////////////////
+	ctx := dlog.WithField(_ctx, "btrfs.inspect.rebuild-trees.read.substep", "read-superblock")
 	dlog.Info(ctx, "Reading superblock...")
 	sb, err := fs.Superblock()
 	if err != nil {
 		return ScanDevicesResult{}, err
 	}
 
-	dlog.Infof(ctx, "Reading node data from FS...")
-
-	var stats textui.Portion[int]
-	stats.D = len(nodeList)
-	progressWriter := textui.NewProgress[textui.Portion[int]](
-		dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.read.substep", "read-nodes"),
-		dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-
+	// read-roots //////////////////////////////////////////////////////////////////
+	ctx = dlog.WithField(_ctx, "btrfs.inspect.rebuild-trees.read.substep", "read-roots")
 	ret := ScanDevicesResult{
 		Superblock: *sb,
 
 		Graph: btrfsutil.NewGraph(ctx, *sb),
-
 		Flags: make(map[btrfsutil.ItemPtr]FlagsAndErr),
 		Names: make(map[btrfsutil.ItemPtr][]byte),
 		Sizes: make(map[btrfsutil.ItemPtr]SizeAndErr),
 	}
 
+	// read-nodes //////////////////////////////////////////////////////////////////
+	ctx = dlog.WithField(_ctx, "btrfs.inspect.rebuild-trees.read.substep", "read-nodes")
+	dlog.Infof(ctx, "Reading node data from FS...")
+	var stats textui.Portion[int]
+	stats.D = len(nodeList)
+	progressWriter := textui.NewProgress[textui.Portion[int]](
+		ctx,
+		dlog.LogLevelInfo,
+		textui.Tunable(1*time.Second))
 	progressWriter.Set(stats)
 	for _, laddr := range nodeList {
 		if err := ctx.Err(); err != nil {
@@ -78,11 +82,8 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalA
 			fs.ReleaseNode(node)
 			return ScanDevicesResult{}, err
 		}
-
 		ret.insertNode(node)
-
 		fs.ReleaseNode(node)
-
 		stats.N++
 		progressWriter.Set(stats)
 	}
@@ -92,7 +93,8 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalA
 	progressWriter.Done()
 	dlog.Info(ctx, "... done reading node data")
 
-	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.read.substep", "check")
+	// check ///////////////////////////////////////////////////////////////////////
+	ctx = dlog.WithField(_ctx, "btrfs.inspect.rebuild-trees.read.substep", "check")
 	if err := ret.Graph.FinalCheck(ctx, fs); err != nil {
 		return ScanDevicesResult{}, err
 	}
diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index fe7fe70..7863e0d 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/datawire/dlib/dlog"
 
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
@@ -283,3 +284,57 @@ func (g Graph) FinalCheck(ctx context.Context, fs btrfstree.NodeSource) error {
 
 	return nil
 }
+
+func ReadGraph(_ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr) (Graph, error) {
+	// read-superblock /////////////////////////////////////////////////////////////
+	ctx := dlog.WithField(_ctx, "btrfs.util.read-graph.step", "read-superblock")
+	dlog.Info(ctx, "Reading superblock...")
+	sb, err := fs.Superblock()
+	if err != nil {
+		return Graph{}, err
+	}
+
+	// read-roots //////////////////////////////////////////////////////////////////
+	ctx = dlog.WithField(_ctx, "btrfs.util.read-graph.step", "read-roots")
+	graph := NewGraph(ctx, *sb)
+
+	// read-nodes //////////////////////////////////////////////////////////////////
+	ctx = dlog.WithField(_ctx, "btrfs.util.read-graph.step", "read-nodes")
+	dlog.Infof(ctx, "Reading node data from FS...")
+	var stats textui.Portion[int]
+	stats.D = len(nodeList)
+	progressWriter := textui.NewProgress[textui.Portion[int]](
+		ctx,
+		dlog.LogLevelInfo,
+		textui.Tunable(1*time.Second))
+	progressWriter.Set(stats)
+	for _, laddr := range nodeList {
+		if err := ctx.Err(); err != nil {
+			return Graph{}, err
+		}
+		node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
+			LAddr: containers.OptionalValue(laddr),
+		})
+		if err != nil {
+			fs.ReleaseNode(node)
+			return Graph{}, err
+		}
+		graph.InsertNode(node)
+		fs.ReleaseNode(node)
+		stats.N++
+		progressWriter.Set(stats)
+	}
+	if stats.N != stats.D {
+		panic("should not happen")
+	}
+	progressWriter.Done()
+	dlog.Info(ctx, "... done reading node data")
+
+	// check ///////////////////////////////////////////////////////////////////////
+	ctx = dlog.WithField(_ctx, "btrfs.util.read-graph.step", "check")
+	if err := graph.FinalCheck(ctx, fs); err != nil {
+		return Graph{}, err
+	}
+
+	return graph, nil
+}
diff --git a/lib/textui/log.go b/lib/textui/log.go
index 9aff364..e8325c2 100644
--- a/lib/textui/log.go
+++ b/lib/textui/log.go
@@ -328,6 +328,10 @@ func fieldOrd(key string) int {
 	case "btrfs.inspect.rebuild-trees.rebuild.want.reason":
 		return -8
 
+	// btrfsutil.Graph /////////////////////////////////////////////////////
+	case "btrfs.util.read-graph.step":
+		return -1
+
 	// btrfsutil.RebuiltForrest ////////////////////////////////////////////
 	case "btrfs.util.rebuilt-forrest.add-tree":
 		return -7
@@ -426,10 +430,14 @@ func writeField(w io.Writer, key string, val any) {
 					name = strings.TrimPrefix(name, "rebuild.")
 				}
 			}
-		case strings.HasPrefix(name, "util.rebuilt-forrest."):
-			name = strings.TrimPrefix(name, "util.rebuilt-forrest.")
-		case strings.HasPrefix(name, "util.rebuilt-tree."):
-			name = strings.TrimPrefix(name, "util.rebuilt-tree.")
+		case strings.HasPrefix(name, "util."):
+			name = strings.TrimPrefix(name, "util.")
+			switch {
+			case strings.HasPrefix(name, "rebuilt-forrest."):
+				name = strings.TrimPrefix(name, "rebuilt-forrest.")
+			case strings.HasPrefix(name, "rebuilt-tree."):
+				name = strings.TrimPrefix(name, "rebuilt-tree.")
+			}
 		}
 	}