btrfsutil: RebuiltForrest: Add a .RebuiltAddRoots() method, cmd/btrfs-rec: Add a --trees flag

4 files changed, 77 insertions, 10 deletions

diff --git a/cmd/btrfs-rec/main.go b/cmd/btrfs-rec/main.go
index e167095..39ba9ef 100644
--- a/cmd/btrfs-rec/main.go
+++ b/cmd/btrfs-rec/main.go
@@ -17,8 +17,10 @@ import (
 	"github.com/spf13/cobra"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"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/btrfsutil"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/profile"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
@@ -49,9 +51,10 @@ var globalFlags struct {
 	logLevel textui.LogLevelFlag
 	pvs      []string
 
-	mappings string
-	nodeList string
-	rebuild  bool
+	mappings  string
+	nodeList  string
+	rebuild   bool
+	treeRoots string
 
 	stopProfiling profile.StopFunc
 
@@ -104,6 +107,10 @@ func main() {
 	argparser.PersistentFlags().BoolVar(&globalFlags.rebuild, "rebuild", false,
 		"attempt to rebuild broken btrees when reading")
 
+	argparser.PersistentFlags().StringVar(&globalFlags.treeRoots, "trees", "",
+		"load list of tree roots (output of 'btrfs-recs inspect rebuild-trees') from external JSON file `trees.json`; implies --rebuild")
+	noError(argparser.MarkPersistentFlagFilename("trees"))
+
 	globalFlags.stopProfiling = profile.AddProfileFlags(argparser.PersistentFlags(), "profile.")
 
 	globalFlags.openFlag = os.O_RDONLY
@@ -210,7 +217,7 @@ func runWithRawFSAndNodeList(runE func(*btrfs.FS, []btrfsvol.LogicalAddr, *cobra
 func _runWithReadableFS(wantNodeList bool, runE func(btrfs.ReadableFS, []btrfsvol.LogicalAddr, *cobra.Command, []string) error) func(*cobra.Command, []string) error {
 	inner := func(fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr, cmd *cobra.Command, args []string) error {
 		var rfs btrfs.ReadableFS = fs
-		if globalFlags.rebuild {
+		if globalFlags.rebuild || globalFlags.treeRoots != "" {
 			ctx := cmd.Context()
 
 			graph, err := btrfsutil.ReadGraph(ctx, fs, nodeList)
@@ -218,14 +225,24 @@ func _runWithReadableFS(wantNodeList bool, runE func(btrfs.ReadableFS, []btrfsvo
 				return err
 			}
 
-			rfs = btrfsutil.NewRebuiltForrest(fs, graph, nil, true)
+			_rfs := btrfsutil.NewRebuiltForrest(fs, graph, nil, true)
+
+			if globalFlags.treeRoots != "" {
+				roots, err := readJSONFile[map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]](ctx, globalFlags.treeRoots)
+				if err != nil {
+					return err
+				}
+				_rfs.RebuiltAddRoots(ctx, roots)
+			}
+
+			rfs = _rfs
 		}
 
 		return runE(rfs, nodeList, cmd, args)
 	}
 
 	return func(cmd *cobra.Command, args []string) error {
-		if wantNodeList || globalFlags.rebuild {
+		if wantNodeList || globalFlags.rebuild || globalFlags.treeRoots != "" {
 			return runWithRawFSAndNodeList(inner)(cmd, args)
 		}
 		return runWithRawFS(func(fs *btrfs.FS, cmd *cobra.Command, args []string) error {
diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index 0ff45a9..4ce2435 100644
--- a/lib/btrfsutil/rebuilt_forrest.go
+++ b/lib/btrfsutil/rebuilt_forrest.go
@@ -26,11 +26,15 @@ import (
 // Additionally, it provides some functionality on top of a vanilla
 // btrfs.ReadableFS:
 //
+//   - it provides a RebuiltTree.RebuiltAddRoot() method for repairing a
+//     tree.
+//
 //   - it provides a RebuiltForrest.RebuiltListRoots() method for
 //     listing how trees have been repaired.
 //
-//   - it provides a RebuiltTree.RebuiltAddRoot() method for repairing a
-//     tree.
+//   - it provides a RebuiltForrest.RebuiltAddRoots() method for
+//     batch-loading the results from
+//     RebuiltForrest.RebuiltListroots().
 //
 //   - it provides several RebuiltTree methods that provide advice on
 //     what roots should be added to a tree in order to repair it:
@@ -260,6 +264,48 @@ func (ts *RebuiltForrest) RebuiltListRoots(ctx context.Context) map[btrfsprim.Ob
 	return ret
 }
 
+// RebuiltAddRoots takes a listing of the root nodes for trees (as
+// returned by RebuiltListRoots), and augments the trees to include
+// them.
+func (ts *RebuiltForrest) RebuiltAddRoots(ctx context.Context, roots map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]) {
+	ctx = ts.treesMu.Lock(ctx)
+	defer ts.treesMu.Unlock()
+
+	essentialTrees := []btrfsprim.ObjID{
+		btrfsprim.ROOT_TREE_OBJECTID,
+		btrfsprim.UUID_TREE_OBJECTID,
+	}
+
+	for _, treeID := range essentialTrees {
+		treeRoots, ok := roots[treeID]
+		if !ok {
+			continue
+		}
+		tree, err := ts.RebuiltTree(ctx, treeID)
+		if err != nil {
+			dlog.Errorf(ctx, "RebuiltForrest.RebuiltAddRoots: cannot load essential tree %v: %v", treeID, err)
+			return
+		}
+		for _, root := range maps.SortedKeys(treeRoots) {
+			tree.RebuiltAddRoot(ctx, root)
+		}
+	}
+
+	for _, treeID := range maps.SortedKeys(roots) {
+		if slices.Contains(treeID, essentialTrees) {
+			continue
+		}
+		tree, err := ts.RebuiltTree(ctx, treeID)
+		if err != nil {
+			dlog.Errorf(ctx, "RebuiltForrest.RebuiltAddRoots: cannot load non-essential tree %v: %v", treeID, err)
+			continue
+		}
+		for _, root := range maps.SortedKeys(roots[treeID]) {
+			tree.RebuiltAddRoot(ctx, root)
+		}
+	}
+}
+
 // btrfs.ReadableFS interface //////////////////////////////////////////////////////////////////////////////////////////
 
 var _ btrfs.ReadableFS = (*RebuiltForrest)(nil)
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index 97308a3..f002ea6 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -700,6 +700,10 @@ func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.L
 	tree.mu.Lock()
 	defer tree.mu.Unlock()
 
+	if tree.Roots.Has(rootNode) {
+		return
+	}
+
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
 	dlog.Info(ctx, "adding root...")
 
diff --git a/scripts/main.sh b/scripts/main.sh
index 4e0596b..082576a 100755
--- a/scripts/main.sh
+++ b/scripts/main.sh
@@ -79,11 +79,11 @@ run-btrfs-rec $gendir/3.trees.json \
 run-btrfs-rec $gendir/4.ls-files.txt \
     --mappings=$gendir/2.mappings.json \
     --node-list=$gendir/0.nodes.json \
-    --rebuild \
+    --trees=$gendir/3.trees.json \
     inspect ls-files
 
 run-btrfs-rec $gendir/4.ls-trees.txt \
     --mappings=$gendir/2.mappings.json \
     --node-list=$gendir/0.nodes.json \
-    --rebuild \
+    --trees=$gendir/3.trees.json \
     inspect ls-trees