btrfsutil: RebuiltTree: Take better advantage of cache pinning

2 files changed, 44 insertions, 23 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
index 2160969..427070a 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
@@ -198,11 +198,13 @@ func (o *rebuilder) processAddedItemQueue(ctx context.Context) error {
 
 		ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.settle.item", key)
 		tree := o.rebuilt.RebuiltTree(ctx, key.TreeID)
-		incPtr, ok := tree.RebuiltItems(ctx).Load(key.Key)
+		incPtr, ok := tree.RebuiltAcquireItems(ctx).Load(key.Key)
+		tree.RebuiltReleaseItems()
 		if !ok {
 			panic(fmt.Errorf("should not happen: failed to load already-added item: %v", key))
 		}
-		excPtr, ok := tree.RebuiltPotentialItems(ctx).Load(key.Key)
+		excPtr, ok := tree.RebuiltAcquirePotentialItems(ctx).Load(key.Key)
+		tree.RebuiltReleasePotentialItems()
 		if ok && tree.RebuiltShouldReplace(incPtr.Node, excPtr.Node) {
 			wantKey := wantWithTree{
 				TreeID: key.TreeID,
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
index eff2a83..0526e93 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
@@ -43,18 +43,22 @@ func (o graphCallbacks) Want(ctx context.Context, reason string, treeID btrfspri
 }
 
 func (o *rebuilder) _want(ctx context.Context, wantKey wantWithTree) (key btrfsprim.Key, ok bool) {
-	if o.rebuilt.RebuiltTree(ctx, wantKey.TreeID) == nil {
+	tree := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)
+	if tree == nil {
 		o.enqueueRetry(wantKey.TreeID)
 		return btrfsprim.Key{}, false
 	}
 
+	tgt := wantKey.Key.Key()
+
 	// check if we already have it
 
-	tgt := wantKey.Key.Key()
-	if key, _, ok := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).RebuiltItems(ctx).Search(func(key btrfsprim.Key, _ btrfsutil.ItemPtr) int {
+	key, _, ok = tree.RebuiltAcquireItems(ctx).Search(func(key btrfsprim.Key, _ btrfsutil.ItemPtr) int {
 		key.Offset = 0
 		return tgt.Compare(key)
-	}); ok {
+	})
+	tree.RebuiltReleaseItems()
+	if ok {
 		return key, true
 	}
 
@@ -64,15 +68,16 @@ func (o *rebuilder) _want(ctx context.Context, wantKey wantWithTree) (key btrfsp
 		return btrfsprim.Key{}, false
 	}
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).RebuiltPotentialItems(ctx).Subrange(
+	tree.RebuiltAcquirePotentialItems(ctx).Subrange(
 		func(k btrfsprim.Key, _ btrfsutil.ItemPtr) int {
 			k.Offset = 0
 			return tgt.Compare(k)
 		},
 		func(_ btrfsprim.Key, v btrfsutil.ItemPtr) bool {
-			wants.InsertFrom(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).RebuiltLeafToRoots(ctx, v.Node))
+			wants.InsertFrom(tree.RebuiltLeafToRoots(ctx, v.Node))
 			return true
 		})
+	tree.RebuiltReleasePotentialItems()
 	o.wantAugment(ctx, wantKey, wants)
 	return btrfsprim.Key{}, false
 }
@@ -93,15 +98,19 @@ func (o graphCallbacks) WantOff(ctx context.Context, reason string, treeID btrfs
 }
 
 func (o *rebuilder) _wantOff(ctx context.Context, wantKey wantWithTree) (ok bool) {
-	if o.rebuilt.RebuiltTree(ctx, wantKey.TreeID) == nil {
+	tree := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)
+	if tree == nil {
 		o.enqueueRetry(wantKey.TreeID)
 		return false
 	}
 
+	tgt := wantKey.Key.Key()
+
 	// check if we already have it
 
-	tgt := wantKey.Key.Key()
-	if _, ok := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).RebuiltItems(ctx).Load(tgt); ok {
+	_, ok = tree.RebuiltAcquireItems(ctx).Load(tgt)
+	tree.RebuiltReleaseItems()
+	if ok {
 		return true
 	}
 
@@ -111,12 +120,13 @@ func (o *rebuilder) _wantOff(ctx context.Context, wantKey wantWithTree) (ok bool
 		return false
 	}
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).RebuiltPotentialItems(ctx).Subrange(
+	tree.RebuiltAcquirePotentialItems(ctx).Subrange(
 		func(k btrfsprim.Key, _ btrfsutil.ItemPtr) int { return tgt.Compare(k) },
 		func(_ btrfsprim.Key, v btrfsutil.ItemPtr) bool {
-			wants.InsertFrom(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).RebuiltLeafToRoots(ctx, v.Node))
+			wants.InsertFrom(tree.RebuiltLeafToRoots(ctx, v.Node))
 			return true
 		})
+	tree.RebuiltReleasePotentialItems()
 	o.wantAugment(ctx, wantKey, wants)
 	return false
 }
@@ -134,16 +144,18 @@ func (o graphCallbacks) WantDirIndex(ctx context.Context, reason string, treeID
 	}
 	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
 
-	if o.rebuilt.RebuiltTree(ctx, treeID) == nil {
+	tree := o.rebuilt.RebuiltTree(ctx, treeID)
+	if tree == nil {
 		o.enqueueRetry(treeID)
 		return
 	}
 
+	tgt := wantKey.Key.Key()
+
 	// check if we already have it
 
-	tgt := wantKey.Key.Key()
 	found := false
-	o.rebuilt.RebuiltTree(ctx, treeID).RebuiltItems(ctx).Subrange(
+	tree.RebuiltAcquireItems(ctx).Subrange(
 		func(key btrfsprim.Key, _ btrfsutil.ItemPtr) int {
 			key.Offset = 0
 			return tgt.Compare(key)
@@ -154,6 +166,7 @@ func (o graphCallbacks) WantDirIndex(ctx context.Context, reason string, treeID
 			}
 			return !found
 		})
+	tree.RebuiltReleaseItems()
 	if found {
 		return
 	}
@@ -164,17 +177,18 @@ func (o graphCallbacks) WantDirIndex(ctx context.Context, reason string, treeID
 		return
 	}
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.RebuiltTree(ctx, treeID).RebuiltPotentialItems(ctx).Subrange(
+	tree.RebuiltAcquirePotentialItems(ctx).Subrange(
 		func(key btrfsprim.Key, _ btrfsutil.ItemPtr) int {
 			key.Offset = 0
 			return tgt.Compare(key)
 		},
 		func(_ btrfsprim.Key, ptr btrfsutil.ItemPtr) bool {
 			if itemName, ok := o.scan.Names[ptr]; ok && bytes.Equal(itemName, name) {
-				wants.InsertFrom(o.rebuilt.RebuiltTree(ctx, treeID).RebuiltLeafToRoots(ctx, ptr.Node))
+				wants.InsertFrom(tree.RebuiltLeafToRoots(ctx, ptr.Node))
 			}
 			return true
 		})
+	tree.RebuiltReleasePotentialItems()
 	o.wantAugment(ctx, wantKey, wants)
 }
 
@@ -259,7 +273,8 @@ func (o graphCallbacks) _wantRange(
 	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
 	wantKey.Key.OffsetType = offsetRange
 
-	if o.rebuilt.RebuiltTree(ctx, treeID) == nil {
+	tree := o.rebuilt.RebuiltTree(ctx, treeID)
+	if tree == nil {
 		o.enqueueRetry(treeID)
 		return
 	}
@@ -275,7 +290,7 @@ func (o graphCallbacks) _wantRange(
 	})
 	o._walkRange(
 		ctx,
-		o.rebuilt.RebuiltTree(ctx, treeID).RebuiltItems(ctx),
+		tree.RebuiltAcquireItems(ctx),
 		treeID, objID, typ, beg, end,
 		func(runKey btrfsprim.Key, _ btrfsutil.ItemPtr, runBeg, runEnd uint64) {
 			var overlappingGaps []*containers.RBNode[gap]
@@ -315,12 +330,13 @@ func (o graphCallbacks) _wantRange(
 				})
 			}
 		})
+	tree.RebuiltReleaseItems()
 
 	// Step 2: Fill each gap.
 	if gaps.Len() == 0 {
 		return
 	}
-	potentialItems := o.rebuilt.RebuiltTree(ctx, treeID).RebuiltPotentialItems(ctx)
+	potentialItems := tree.RebuiltAcquirePotentialItems(ctx)
 	gaps.Range(func(rbNode *containers.RBNode[gap]) bool {
 		gap := rbNode.Value
 		last := gap.Beg
@@ -340,7 +356,7 @@ func (o graphCallbacks) _wantRange(
 				wantKey.Key.OffsetLow = gap.Beg
 				wantKey.Key.OffsetHigh = gap.End
 				wantCtx := withWant(ctx, logFieldItemWant, reason, wantKey)
-				o.wantAugment(wantCtx, wantKey, o.rebuilt.RebuiltTree(wantCtx, treeID).RebuiltLeafToRoots(wantCtx, v.Node))
+				o.wantAugment(wantCtx, wantKey, tree.RebuiltLeafToRoots(wantCtx, v.Node))
 				last = runEnd
 			})
 		if last < gap.End {
@@ -352,6 +368,7 @@ func (o graphCallbacks) _wantRange(
 		}
 		return true
 	})
+	tree.RebuiltReleasePotentialItems()
 }
 
 // WantCSum implements btrfscheck.GraphCallbacks.
@@ -372,7 +389,9 @@ func (o graphCallbacks) WantCSum(ctx context.Context, reason string, inodeTree,
 		o.enqueueRetry(inodeTree)
 		return
 	}
-	inodePtr, ok := o.rebuilt.RebuiltTree(inodeCtx, inodeTree).RebuiltItems(inodeCtx).Load(inodeWant.Key.Key())
+	tree := o.rebuilt.RebuiltTree(inodeCtx, inodeTree)
+	inodePtr, ok := tree.RebuiltAcquireItems(inodeCtx).Load(inodeWant.Key.Key())
+	tree.RebuiltReleaseItems()
 	if !ok {
 		panic(fmt.Errorf("should not happen: could not load key: %v", inodeWant))
 	}