btrfsutil: RebuiltForrest: Be lazier about calling .RebuiltAddRoot()

2 files changed, 21 insertions, 10 deletions

diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index beb7d40..019d824 100644
--- a/lib/btrfsutil/rebuilt_forrest.go
+++ b/lib/btrfsutil/rebuilt_forrest.go
@@ -124,6 +124,7 @@ func (ts *RebuiltForrest) RebuiltTree(ctx context.Context, treeID btrfsprim.ObjI
 	if tree.rootErr != nil {
 		return nil, tree.rootErr
 	}
+	tree.initRoots(ctx)
 	return tree, nil
 }
 
@@ -156,27 +157,26 @@ func (ts *RebuiltForrest) rebuildTree(ctx context.Context, treeID btrfsprim.ObjI
 		Roots:   make(containers.Set[btrfsvol.LogicalAddr]),
 		forrest: ts,
 	}
-	var root btrfsvol.LogicalAddr
 	switch treeID {
 	case btrfsprim.ROOT_TREE_OBJECTID:
 		sb, _ := ts.Superblock()
-		root = sb.RootTree
+		ts.trees[treeID].Root = sb.RootTree
 	case btrfsprim.CHUNK_TREE_OBJECTID:
 		sb, _ := ts.Superblock()
-		root = sb.ChunkTree
+		ts.trees[treeID].Root = sb.ChunkTree
 	case btrfsprim.TREE_LOG_OBJECTID:
 		sb, _ := ts.Superblock()
-		root = sb.LogTree
+		ts.trees[treeID].Root = sb.LogTree
 	case btrfsprim.BLOCK_GROUP_TREE_OBJECTID:
 		sb, _ := ts.Superblock()
-		root = sb.BlockGroupRoot
+		ts.trees[treeID].Root = sb.BlockGroupRoot
 	default:
 		rootOff, rootItem, ok := ts.cb.LookupRoot(ctx, treeID)
 		if !ok {
 			ts.trees[treeID].rootErr = btrfstree.ErrNoTree
 			return
 		}
-		root = rootItem.ByteNr
+		ts.trees[treeID].Root = rootItem.ByteNr
 		ts.trees[treeID].UUID = rootItem.UUID
 		if rootItem.ParentUUID != (btrfsprim.UUID{}) {
 			ts.trees[treeID].ParentGen = rootOff
@@ -197,10 +197,6 @@ func (ts *RebuiltForrest) rebuildTree(ctx context.Context, treeID btrfsprim.ObjI
 			}
 		}
 	}
-
-	if root != 0 {
-		ts.trees[treeID].RebuiltAddRoot(ctx, root)
-	}
 }
 
 func (ts *RebuiltForrest) flushNegativeCache(ctx context.Context) {
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index fc36ebb..d890dc4 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -30,12 +30,15 @@ type RebuiltTree struct {
 
 	ID        btrfsprim.ObjID
 	UUID      btrfsprim.UUID
+	Root      btrfsvol.LogicalAddr
 	Parent    *RebuiltTree
 	ParentGen btrfsprim.Generation // offset of this tree's root item
 	forrest   *RebuiltForrest
 
 	// mutable
 
+	initRootsOnce sync.Once
+
 	mu sync.RWMutex
 
 	Roots containers.Set[btrfsvol.LogicalAddr]
@@ -79,6 +82,14 @@ func makeRebuiltSharedCache(forrest *RebuiltForrest) rebuiltSharedCache {
 	return ret
 }
 
+func (tree *RebuiltTree) initRoots(ctx context.Context) {
+	tree.initRootsOnce.Do(func() {
+		if tree.Root != 0 {
+			tree.RebuiltAddRoot(ctx, tree.Root)
+		}
+	})
+}
+
 // evictable member 1: .acquireNodeIndex() /////////////////////////////////////////////////////////////////////////////
 
 type rebuiltRoots = map[btrfsvol.LogicalAddr]rebuiltPathInfo
@@ -273,6 +284,7 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 //
 // When done with the map, call .RebuiltReleaseItems().
 func (tree *RebuiltTree) RebuiltAcquireItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
 
@@ -293,6 +305,7 @@ func (tree *RebuiltTree) RebuiltReleaseItems() {
 //
 // When done with the map, call .RebuiltReleasePotentialItems().
 func (tree *RebuiltTree) RebuiltAcquirePotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
 
@@ -490,6 +503,7 @@ func (tree *RebuiltTree) RebuiltLeafToRoots(ctx context.Context, leaf btrfsvol.L
 			tree.ID, leaf))
 	}
 
+	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
 
@@ -589,6 +603,7 @@ func (tree *RebuiltTree) TreeSubrange(ctx context.Context,
 
 // TreeWalk implements btrfstree.Tree.
 func (tree *RebuiltTree) TreeWalk(ctx context.Context, cbs btrfstree.TreeWalkHandler) {
+	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()