1 files changed, 49 insertions, 18 deletions

diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index 79d8beb..4ce2435 100644
--- a/lib/btrfsutil/rebuilt_forrest.go
+++ b/lib/btrfsutil/rebuilt_forrest.go
@@ -23,29 +23,18 @@ import (
 // RebuiltForrest is an abstraction for rebuilding and accessing
 // potentially broken btrees.
 //
-// It is conceptually a btrfstree.Forrest, and adds similar
-// broken-tree handling to OldRebuiltForrest.  However, it is much
-// more efficient than OldRebuiltForrest.
+// Additionally, it provides some functionality on top of a vanilla
+// btrfs.ReadableFS:
 //
-// The efficiency improvements are possible because of the API
-// differences, which are necessary for how it is used in
-// rebuildtrees:
-//
-//   - it consumes an already-read Graph instead of reading the graph
-//     itself
-//
-//   - it does not use `btrfstree.Path`
-//
-//   - it does not keep track of errors encountered in a tree
-//
-// Additionally, it provides some functionality that OldRebuiltForrest
-// does not:
+//   - 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:
@@ -275,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)