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

2 files changed, 52 insertions, 2 deletions

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