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

2 files changed, 67 insertions, 4 deletions

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.")
+			}
 		}
 	}