rebuildnodes: Optimize: Try to avoid disk access for DIR_INDEX

4 files changed, 39 insertions, 25 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index b53a28e..b1ae7be 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -363,15 +363,23 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd
 	}
 }
 
+// Resolve a key to a keyio.ItemPtr.
+//
+// It is not nescessary to call AddTree for that tree first; Resolve will
+// call it for you.
+func (ts *RebuiltTrees) Resolve(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
+	if !ts.AddTree(ctx, treeID) {
+		return keyio.ItemPtr{}, false
+	}
+	return ts.trees[treeID].Items.Load(key)
+}
+
 // Load reads an item from a tree.
 //
 // It is not nescessary to call AddTree for that tree first; Load will
 // call it for you.
 func (ts *RebuiltTrees) Load(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
-	if !ts.AddTree(ctx, treeID) {
-		return nil, false
-	}
-	ptr, ok := ts.trees[treeID].Items.Load(key)
+	ptr, ok := ts.Resolve(ctx, treeID, key)
 	if !ok {
 		return nil, false
 	}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index 24c3dcf..eb03004 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -39,6 +39,7 @@ type Handle struct {
 	sb      btrfstree.Superblock
 	graph   graph.Graph
 
+	Names map[ItemPtr][]byte
 	Sizes map[ItemPtr]SizeAndErr
 
 	cache *containers.LRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
@@ -49,6 +50,7 @@ func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock)
 		rawFile: file,
 		sb:      sb,
 
+		Names: make(map[ItemPtr][]byte),
 		Sizes: make(map[ItemPtr]SizeAndErr),
 
 		cache: containers.NewLRUCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]](textui.Tunable(8)),
@@ -62,6 +64,10 @@ func (o *Handle) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.
 			Idx:  i,
 		}
 		switch itemBody := item.Body.(type) {
+		case btrfsitem.DirEntry:
+			if item.Key.ItemType == btrfsprim.DIR_INDEX_KEY {
+				o.Names[ptr] = append([]byte(nil), itemBody.Name...)
+			}
 		case btrfsitem.ExtentCSum:
 			o.Sizes[ptr] = SizeAndErr{
 				Size: uint64(itemBody.Size()),
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 7e55732..4f5990e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -5,6 +5,7 @@
 package rebuildnodes
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"sort"
@@ -456,12 +457,7 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, tgt bt
 	return false
 }
 
-// wantFunc implements rebuildCallbacks.
-func (o *rebuilder) wantFunc(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key",
-		fmt.Sprintf("tree=%v key={%v %v ?} +func", treeID, objID, typ))
-
+func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(keyio.ItemPtr) bool) {
 	if !o.rebuilt.AddTree(ctx, treeID) {
 		o.itemQueue = append(o.itemQueue, o.curKey)
 		return
@@ -478,11 +474,11 @@ func (o *rebuilder) wantFunc(ctx context.Context, reason string, treeID btrfspri
 		return tgt.Cmp(key)
 	})
 	for _, itemKey := range keys {
-		itemBody, ok := o.rebuilt.Load(ctx, treeID, itemKey)
+		itemPtr, ok := o.rebuilt.Resolve(ctx, treeID, itemKey)
 		if !ok {
-			o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", itemKey))
+			o.ioErr(ctx, fmt.Errorf("could not resolve previously read item: %v", itemKey))
 		}
-		if fn(itemBody) {
+		if fn(itemPtr) {
 			return
 		}
 	}
@@ -493,11 +489,7 @@ func (o *rebuilder) wantFunc(ctx context.Context, reason string, treeID btrfspri
 	o.rebuilt.Keys(treeID).Subrange(
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
 		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
-			itemBody, ok := o.keyIO.ReadItem(ctx, v)
-			if !ok {
-				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v at %v", k, v))
-			}
-			if fn(itemBody) {
+			if fn(v) {
 				wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			}
 			return true
@@ -505,6 +497,18 @@ func (o *rebuilder) wantFunc(ctx context.Context, reason string, treeID btrfspri
 	o.wantAugment(ctx, treeID, wants)
 }
 
+// wantDirIndex implements rebuildCallbacks.
+func (o *rebuilder) wantDirIndex(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, name []byte) {
+	typ := btrfsitem.DIR_INDEX_KEY
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key",
+		fmt.Sprintf("tree=%v key={%v %v ?} name=%q", treeID, objID, typ, name))
+	o._wantFunc(ctx, treeID, objID, typ, func(ptr keyio.ItemPtr) bool {
+		itemName, ok := o.keyIO.Names[ptr]
+		return ok && bytes.Equal(itemName, name)
+	})
+}
+
 func (o *rebuilder) _wantRange(
 	ctx context.Context,
 	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go
index bf4c95d..271595e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_graph.go
@@ -7,7 +7,6 @@ package rebuildnodes
 import (
 	"context"
 	"fmt"
-	"reflect"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
@@ -20,7 +19,7 @@ type rebuildCallbacks interface {
 	fsErr(ctx context.Context, e error)
 	want(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType)
 	wantOff(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64)
-	wantFunc(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool)
+	wantDirIndex(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, name []byte)
 	wantCSum(ctx context.Context, reason string, beg, end btrfsvol.LogicalAddr) // interval is [beg, end)
 	wantFileExt(ctx context.Context, reason string, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64)
 }
@@ -65,13 +64,10 @@ func handleItem(o rebuildCallbacks, ctx context.Context, treeID btrfsprim.ObjID,
 		// siblings
 		switch item.Key.ItemType {
 		case btrfsitem.DIR_ITEM_KEY:
-			o.wantFunc(ctx, "corresponding DIR_INDEX",
+			o.wantDirIndex(ctx, "corresponding DIR_INDEX",
 				treeID,
 				item.Key.ObjectID,
-				btrfsitem.DIR_INDEX_KEY,
-				func(peerItem btrfsitem.Item) bool {
-					return reflect.DeepEqual(item, peerItem)
-				})
+				body.Name)
 		case btrfsitem.DIR_INDEX_KEY:
 			o.wantOff(ctx, "corresponding DIR_ITEM",
 				treeID,