Merge branch 'lukeshu/rebuilt-misc'

6 files changed, 85 insertions, 87 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
index 7e2e49f..77d56ae 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
@@ -159,7 +159,7 @@ func (o *rebuilder) processTreeQueue(ctx context.Context) error {
 		}
 		// This will call o.AddedItem as nescessary, which
 		// inserts to o.addedItemQueue.
-		_, _ = o.rebuilt.RebuiltTree(ctx, o.curKey.TreeID)
+		_, _ = o.rebuilt.ForrestLookup(ctx, o.curKey.TreeID)
 	}
 
 	return nil
@@ -415,7 +415,7 @@ func (o *rebuilder) processSettledItemQueue(ctx context.Context) error {
 			ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.process.item", key)
 			item := keyAndBody{
 				itemToVisit: key,
-				Body:        discardErr(o.rebuilt.RebuiltTree(ctx, key.TreeID)).ReadItem(ctx, key.Key),
+				Body:        discardErr(discardErr(o.rebuilt.RebuiltTree(ctx, key.TreeID)).TreeLookup(ctx, key.Key)).Body,
 			}
 			if key.TreeID == btrfsprim.EXTENT_TREE_OBJECTID &&
 				(key.ItemType == btrfsprim.EXTENT_ITEM_KEY || key.ItemType == btrfsprim.METADATA_ITEM_KEY) {
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go
index e227cc7..15ad42c 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go
@@ -36,7 +36,7 @@ func (o forrestCallbacks) AddedRoot(_ context.Context, tree btrfsprim.ObjID, _ b
 }
 
 // LookupRoot implements btrfsutil.RebuiltForrestCallbacks.
-func (o forrestCallbacks) LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool) {
+func (o forrestCallbacks) LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, _item btrfsitem.Root, ok bool) {
 	wantKey := wantWithTree{
 		TreeID: btrfsprim.ROOT_TREE_OBJECTID,
 		Key: want{
@@ -51,9 +51,9 @@ func (o forrestCallbacks) LookupRoot(ctx context.Context, tree btrfsprim.ObjID)
 		o.enqueueRetry(btrfsprim.ROOT_TREE_OBJECTID)
 		return 0, btrfsitem.Root{}, false
 	}
-	itemBody := discardErr(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)).ReadItem(ctx, foundKey)
-	defer itemBody.Free()
-	switch itemBody := itemBody.(type) {
+	item, _ := discardErr(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)).TreeLookup(ctx, foundKey)
+	defer item.Body.Free()
+	switch itemBody := item.Body.(type) {
 	case *btrfsitem.Root:
 		return btrfsprim.Generation(foundKey.Offset), *itemBody, true
 	case *btrfsitem.Error:
@@ -77,9 +77,9 @@ func (o forrestCallbacks) LookupUUID(ctx context.Context, uuid btrfsprim.UUID) (
 		o.enqueueRetry(btrfsprim.UUID_TREE_OBJECTID)
 		return 0, false
 	}
-	itemBody := discardErr(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)).ReadItem(ctx, wantKey.Key.Key())
-	defer itemBody.Free()
-	switch itemBody := itemBody.(type) {
+	item, _ := discardErr(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)).TreeLookup(ctx, wantKey.Key.Key())
+	defer item.Body.Free()
+	switch itemBody := item.Body.(type) {
 	case *btrfsitem.UUIDMap:
 		return itemBody.ObjID, true
 	case *btrfsitem.Error:
diff --git a/lib/btrfsutil/rebuilt_callbacks.go b/lib/btrfsutil/rebuilt_callbacks.go
index fdc826c..0b51d08 100644
--- a/lib/btrfsutil/rebuilt_callbacks.go
+++ b/lib/btrfsutil/rebuilt_callbacks.go
@@ -10,6 +10,7 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
@@ -25,34 +26,25 @@ type RebuiltForrestExtendedCallbacks interface {
 }
 
 type noopRebuiltForrestCallbacks struct {
-	forrest *RebuiltForrest
+	forrest btrfstree.Forrest
 }
 
 func (noopRebuiltForrestCallbacks) AddedRoot(context.Context, btrfsprim.ObjID, btrfsvol.LogicalAddr) {
 }
 
 func (cb noopRebuiltForrestCallbacks) LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, _item btrfsitem.Root, ok bool) {
-	rootTree, err := cb.forrest.RebuiltTree(ctx, btrfsprim.ROOT_TREE_OBJECTID)
+	rootTree, err := cb.forrest.ForrestLookup(ctx, btrfsprim.ROOT_TREE_OBJECTID)
 	if err != nil {
 		return 0, btrfsitem.Root{}, false
 	}
-	tgt := btrfsprim.Key{
-		ObjectID: tree,
-		ItemType: btrfsprim.ROOT_ITEM_KEY,
-	}
-	itemKey, itemPtr, ok := rootTree.RebuiltAcquireItems(ctx).Search(func(key btrfsprim.Key, _ ItemPtr) int {
-		key.Offset = 0
-		return tgt.Compare(key)
-	})
-	rootTree.RebuiltReleaseItems()
-	if !ok {
+	item, err := rootTree.TreeSearch(ctx, btrfstree.SearchRootItem(tree))
+	if err != nil {
 		return 0, btrfsitem.Root{}, false
 	}
-	item := cb.forrest.readItem(ctx, itemPtr)
 	defer item.Body.Free()
 	switch itemBody := item.Body.(type) {
 	case *btrfsitem.Root:
-		return btrfsprim.Generation(itemKey.Offset), *itemBody, true
+		return btrfsprim.Generation(item.Key.Offset), *itemBody, true
 	case *btrfsitem.Error:
 		return 0, btrfsitem.Root{}, false
 	default:
@@ -63,17 +55,15 @@ func (cb noopRebuiltForrestCallbacks) LookupRoot(ctx context.Context, tree btrfs
 }
 
 func (cb noopRebuiltForrestCallbacks) LookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool) {
-	uuidTree, err := cb.forrest.RebuiltTree(ctx, btrfsprim.UUID_TREE_OBJECTID)
+	uuidTree, err := cb.forrest.ForrestLookup(ctx, btrfsprim.UUID_TREE_OBJECTID)
 	if err != nil {
 		return 0, false
 	}
 	tgt := btrfsitem.UUIDToKey(uuid)
-	itemPtr, ok := uuidTree.RebuiltAcquireItems(ctx).Load(tgt)
-	uuidTree.RebuiltReleaseItems()
-	if !ok {
+	item, err := uuidTree.TreeLookup(ctx, tgt)
+	if err != nil {
 		return 0, false
 	}
-	item := cb.forrest.readItem(ctx, itemPtr)
 	defer item.Body.Free()
 	switch itemBody := item.Body.(type) {
 	case *btrfsitem.UUIDMap:
diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index b935d85..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.inner.Superblock()
-		root = sb.RootTree
+		sb, _ := ts.Superblock()
+		ts.trees[treeID].Root = sb.RootTree
 	case btrfsprim.CHUNK_TREE_OBJECTID:
-		sb, _ := ts.inner.Superblock()
-		root = sb.ChunkTree
+		sb, _ := ts.Superblock()
+		ts.trees[treeID].Root = sb.ChunkTree
 	case btrfsprim.TREE_LOG_OBJECTID:
-		sb, _ := ts.inner.Superblock()
-		root = sb.LogTree
+		sb, _ := ts.Superblock()
+		ts.trees[treeID].Root = sb.LogTree
 	case btrfsprim.BLOCK_GROUP_TREE_OBJECTID:
-		sb, _ := ts.inner.Superblock()
-		root = sb.BlockGroupRoot
+		sb, _ := ts.Superblock()
+		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_readitem.go b/lib/btrfsutil/rebuilt_readitem.go
index e5faa45..841b04a 100644
--- a/lib/btrfsutil/rebuilt_readitem.go
+++ b/lib/btrfsutil/rebuilt_readitem.go
@@ -35,7 +35,7 @@ func (ts *RebuiltForrest) readItem(ctx context.Context, ptr ItemPtr) btrfstree.I
 		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for negative item slot: %v", ptr.Slot))
 	}
 
-	node, err := ts.inner.AcquireNode(ctx, ptr.Node, btrfstree.NodeExpectations{
+	node, err := ts.AcquireNode(ctx, ptr.Node, btrfstree.NodeExpectations{
 		LAddr:      containers.OptionalValue(ptr.Node),
 		Level:      containers.OptionalValue(graphInfo.Level),
 		Generation: containers.OptionalValue(graphInfo.Generation),
@@ -50,13 +50,12 @@ func (ts *RebuiltForrest) readItem(ctx context.Context, ptr ItemPtr) btrfstree.I
 		MinItem: containers.OptionalValue(graphInfo.MinItem(ts.graph)),
 		MaxItem: containers.OptionalValue(graphInfo.MaxItem(ts.graph)),
 	})
-	defer ts.inner.ReleaseNode(node)
+	defer ts.ReleaseNode(node)
 	if err != nil {
 		panic(fmt.Errorf("should not happen: i/o error: %w", err))
 	}
 
 	items := node.BodyLeaf
-
 	if ptr.Slot >= len(items) {
 		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for out-of-bounds item slot: slot=%v len=%v",
 			ptr.Slot, len(items)))
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index 9ab0356..8d5921b 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
@@ -205,33 +216,31 @@ nextKP:
 		}
 
 		indexer.node(ctx, kp.FromNode, stack)
-		if len(indexer.nodeToRoots[kp.FromNode]) > 0 {
-			for root, rootInfo := range indexer.nodeToRoots[kp.FromNode] {
-				if kp.FromSlot+1 < len(indexer.tree.forrest.graph.EdgesFrom[kp.FromNode]) {
-					rootInfo.loMaxItem = indexer.tree.forrest.graph.EdgesFrom[kp.FromNode][kp.FromSlot+1].ToKey.Mm()
-					rootInfo.hiMaxItem = rootInfo.loMaxItem
-				} else {
-					// OK, now check the MaxItem expectation.
-					//
-					// We'll use the hiMaxItem, because it only needs to be
-					// valid in *one* of the paths to this node.
-					kpMaxItem := rootInfo.hiMaxItem
-					if kpMaxItem.Compare(nodeInfo.MaxItem(indexer.tree.forrest.graph)) < 0 {
-						continue nextKP
-					}
-				}
-				oldRootInfo, ok := roots[root]
-				if ok && rootInfo.loMaxItem.Compare(oldRootInfo.loMaxItem) > 0 {
-					rootInfo.loMaxItem = oldRootInfo.loMaxItem
-				}
-				if ok && rootInfo.hiMaxItem.Compare(oldRootInfo.hiMaxItem) < 0 {
-					rootInfo.hiMaxItem = oldRootInfo.hiMaxItem
-				}
-				if roots == nil {
-					roots = make(rebuiltRoots)
+		for root, rootInfo := range indexer.nodeToRoots[kp.FromNode] {
+			if kp.FromSlot+1 < len(indexer.tree.forrest.graph.EdgesFrom[kp.FromNode]) {
+				rootInfo.loMaxItem = indexer.tree.forrest.graph.EdgesFrom[kp.FromNode][kp.FromSlot+1].ToKey.Mm()
+				rootInfo.hiMaxItem = rootInfo.loMaxItem
+			} else {
+				// OK, now check the MaxItem expectation.
+				//
+				// We'll use the hiMaxItem, because it only needs to be
+				// valid in *one* of the paths to this node.
+				kpMaxItem := rootInfo.hiMaxItem
+				if kpMaxItem.Compare(nodeInfo.MaxItem(indexer.tree.forrest.graph)) < 0 {
+					continue nextKP
 				}
-				roots[root] = rootInfo
 			}
+			oldRootInfo, ok := roots[root]
+			if ok && rootInfo.loMaxItem.Compare(oldRootInfo.loMaxItem) > 0 {
+				rootInfo.loMaxItem = oldRootInfo.loMaxItem
+			}
+			if ok && rootInfo.hiMaxItem.Compare(oldRootInfo.hiMaxItem) < 0 {
+				rootInfo.hiMaxItem = oldRootInfo.hiMaxItem
+			}
+			if roots == nil {
+				roots = make(rebuiltRoots)
+			}
+			roots[root] = rootInfo
 		}
 	}
 	if roots == nil {
@@ -273,6 +282,10 @@ 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()
+
 	return tree.forrest.incItems.Acquire(ctx, tree.ID)
 }
 
@@ -290,6 +303,10 @@ 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()
+
 	return tree.forrest.excItems.Acquire(ctx, tree.ID)
 }
 
@@ -301,12 +318,12 @@ func (tree *RebuiltTree) RebuiltReleasePotentialItems() {
 
 func (tree *RebuiltTree) uncachedIncItems(ctx context.Context) containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-inc-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, true)
+	return tree.uncachedItems(ctx, true)
 }
 
 func (tree *RebuiltTree) uncachedExcItems(ctx context.Context) containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-exc-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, false)
+	return tree.uncachedItems(ctx, false)
 }
 
 type rebuiltItemStats struct {
@@ -320,10 +337,7 @@ func (s rebuiltItemStats) String() string {
 		s.Leafs, s.NumItems, s.NumDups)
 }
 
-func (tree *RebuiltTree) items(ctx context.Context, inc bool) containers.SortedMap[btrfsprim.Key, ItemPtr] {
-	tree.mu.RLock()
-	defer tree.mu.RUnlock()
-
+func (tree *RebuiltTree) uncachedItems(ctx context.Context, inc bool) containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	var leafs []btrfsvol.LogicalAddr
 	for node, roots := range tree.acquireNodeIndex(ctx).nodeToRoots {
 		if tree.forrest.graph.Nodes[node].Level == 0 && maps.HaveAnyKeysInCommon(tree.Roots, roots) == inc {
@@ -416,6 +430,7 @@ func (s rebuiltRootStats) String() string {
 func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
 	tree.mu.Lock()
 	defer tree.mu.Unlock()
+
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
 	dlog.Info(ctx, "adding root...")
 
@@ -478,16 +493,6 @@ func (tree *RebuiltTree) RebuiltCOWDistance(parentID btrfsprim.ObjID) (dist int,
 	}
 }
 
-// ReadItem reads an item from a tree.
-func (tree *RebuiltTree) ReadItem(ctx context.Context, key btrfsprim.Key) btrfsitem.Item {
-	ptr, ok := tree.RebuiltAcquireItems(ctx).Load(key)
-	tree.RebuiltReleaseItems()
-	if !ok {
-		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltTree.ReadItem called for not-included key: %v", key))
-	}
-	return tree.forrest.readItem(ctx, ptr).Body
-}
-
 // RebuiltLeafToRoots returns the list of potential roots (to pass to
 // .RebuiltAddRoot) that include a given leaf-node.
 func (tree *RebuiltTree) RebuiltLeafToRoots(ctx context.Context, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
@@ -495,8 +500,11 @@ func (tree *RebuiltTree) RebuiltLeafToRoots(ctx context.Context, leaf btrfsvol.L
 		panic(fmt.Errorf("should not happen: (tree=%v).RebuiltLeafToRoots(leaf=%v): not a leaf",
 			tree.ID, leaf))
 	}
+
+	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
+
 	ret := make(containers.Set[btrfsvol.LogicalAddr])
 	for root := range tree.acquireNodeIndex(ctx).nodeToRoots[leaf] {
 		if tree.Roots.Has(root) {
@@ -593,6 +601,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()
 
@@ -651,7 +660,7 @@ func (walker *rebuiltWalker) walk(ctx context.Context, path btrfstree.Path) {
 	}
 
 	// 001
-	nodeAddr, nodeExp, ok := path.NodeExpectations(ctx, true)
+	nodeAddr, nodeExp, ok := path.NodeExpectations(ctx, false)
 	if !ok {
 		panic(fmt.Errorf("should not happen: btrfsutil.rebuiltWalker.walk called with non-node path: %v",
 			path))
@@ -677,6 +686,10 @@ func (walker *rebuiltWalker) walk(ctx context.Context, path btrfstree.Path) {
 		}
 		return
 	}
+	if !maps.HaveAnyKeysInCommon(walker.tree.Roots, walker.nodeIndex.nodeToRoots[nodeAddr]) {
+		panic(fmt.Errorf("should not happen: node@%v is not in the tree, but did not error: path=%#v exp=%#v",
+			nodeAddr, path, nodeExp))
+	}
 	if walker.visited.Has(nodeAddr) {
 		return
 	}