btrfstree: Change the NodeSource API so that it can do caching

10 files changed, 78 insertions, 49 deletions

diff --git a/lib/btrfs/btrfstree/btree.go b/lib/btrfs/btrfstree/btree.go
index dde485a..4f5d21b 100644
--- a/lib/btrfs/btrfstree/btree.go
+++ b/lib/btrfs/btrfstree/btree.go
@@ -97,5 +97,6 @@ func (e *TreeError) Error() string {
 
 type NodeSource interface {
 	Superblock() (*Superblock, error)
-	ReadNode(ctx context.Context, addr btrfsvol.LogicalAddr, exp NodeExpectations) (*Node, error)
+	AcquireNode(ctx context.Context, addr btrfsvol.LogicalAddr, exp NodeExpectations) (*Node, error)
+	ReleaseNode(*Node)
 }
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index a6311f9..86eab11 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -46,8 +46,8 @@ func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
 	if !ok {
 		return
 	}
-	node, err := tree.Forrest.NodeSource.ReadNode(ctx, nodeAddr, nodeExp)
-	defer node.Free()
+	node, err := tree.Forrest.NodeSource.AcquireNode(ctx, nodeAddr, nodeExp)
+	defer tree.Forrest.NodeSource.ReleaseNode(node)
 	if ctx.Err() != nil {
 		return
 	}
diff --git a/lib/btrfs/btrfstree/types_node.go b/lib/btrfs/btrfstree/types_node.go
index 082be2c..e29431f 100644
--- a/lib/btrfs/btrfstree/types_node.go
+++ b/lib/btrfs/btrfstree/types_node.go
@@ -308,7 +308,9 @@ type ItemHeader struct {
 
 var itemPool containers.SlicePool[Item]
 
-func (node *Node) Free() {
+// RawFree is for low-level use by caches; don't use .RawFree, use
+// ReleaseNode.
+func (node *Node) RawFree() {
 	if node == nil {
 		return
 	}
diff --git a/lib/btrfs/io2_lv.go b/lib/btrfs/io2_lv.go
index 03b2107..72e97f3 100644
--- a/lib/btrfs/io2_lv.go
+++ b/lib/btrfs/io2_lv.go
@@ -16,6 +16,7 @@ import (
 	"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"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 )
 
@@ -27,6 +28,8 @@ type FS struct {
 	cacheSuperblocks []*diskio.Ref[btrfsvol.PhysicalAddr, btrfstree.Superblock]
 	cacheSuperblock  *btrfstree.Superblock
 
+	cacheNodes containers.Cache[btrfsvol.LogicalAddr, nodeCacheEntry]
+
 	cacheObjID2All  map[btrfsprim.ObjID]treeInfo
 	cacheUUID2ObjID map[btrfsprim.UUID]btrfsprim.ObjID
 }
diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index 4d1ebe3..8aa485f 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -12,6 +12,8 @@ import (
 	"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"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 // This file is ordered from low-level to high-level.
@@ -80,23 +82,56 @@ var _ btrfstree.NodeFile = (*FS)(nil)
 
 // btrfstree.NodeSource ////////////////////////////////////////////////////////
 
-// ReadNode implements btrfstree.NodeSource.
-func (fs *FS) ReadNode(_ context.Context, addr btrfsvol.LogicalAddr, exp btrfstree.NodeExpectations) (*btrfstree.Node, error) {
-	sb, err := fs.Superblock()
-	if err != nil {
-		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
+type nodeCacheEntry struct {
+	node *btrfstree.Node
+	err  error
+}
+
+// AcquireNode implements btrfstree.NodeSource.
+func (fs *FS) AcquireNode(ctx context.Context, addr btrfsvol.LogicalAddr, exp btrfstree.NodeExpectations) (*btrfstree.Node, error) {
+	if fs.cacheNodes == nil {
+		fs.cacheNodes = containers.NewARCache[btrfsvol.LogicalAddr, nodeCacheEntry](
+			textui.Tunable(4*(btrfstree.MaxLevel+1)),
+			containers.SourceFunc[btrfsvol.LogicalAddr, nodeCacheEntry](fs.readNode),
+		)
 	}
 
-	node, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, addr)
-	if err != nil {
-		return node, err
+	nodeEntry := fs.cacheNodes.Acquire(ctx, addr)
+	if nodeEntry.err != nil {
+		err := nodeEntry.err
+		fs.cacheNodes.Release(addr)
+		return nil, err
 	}
 
-	if err := exp.Check(node); err != nil {
-		return node, fmt.Errorf("btrfstree.ReadNode: node@%v: %w", addr, err) // fmt.Errorf("btrfs.FS.ReadNode: node@%v: %w", addr, err)
+	if nodeEntry.node != nil {
+		if err := exp.Check(nodeEntry.node); err != nil {
+			fs.cacheNodes.Release(addr)
+			return nil, fmt.Errorf("btrfstree.ReadNode: node@%v: %w", addr, err) // fmt.Errorf("btrfs.FS.AcquireNode: node@%v: %w", addr, err)
+		}
+	}
+
+	return nodeEntry.node, nil
+}
+
+// ReleaseNode implements btrfstree.NodeSource.
+func (fs *FS) ReleaseNode(node *btrfstree.Node) {
+	if node == nil {
+		return
+	}
+	fs.cacheNodes.Release(node.Head.Addr)
+}
+
+func (fs *FS) readNode(_ context.Context, addr btrfsvol.LogicalAddr, nodeEntry *nodeCacheEntry) {
+	nodeEntry.node.RawFree()
+	nodeEntry.node = nil
+
+	sb, err := fs.Superblock()
+	if err != nil {
+		nodeEntry.err = err
+		return
 	}
 
-	return node, nil
+	nodeEntry.node, nodeEntry.err = btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, addr)
 }
 
 var _ btrfstree.NodeSource = (*FS)(nil)
diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index 4148f89..ea33be7 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -207,9 +207,10 @@ func (g Graph) FinalCheck(ctx context.Context, fs btrfstree.NodeSource) error {
 	progressWriter.Set(stats)
 	for laddr := range g.EdgesTo {
 		if _, ok := g.Nodes[laddr]; !ok {
-			_, err := fs.ReadNode(ctx, laddr, btrfstree.NodeExpectations{
+			node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
 				LAddr: containers.OptionalValue(laddr),
 			})
+			fs.ReleaseNode(node)
 			if err == nil {
 				progressWriter.Done()
 				return fmt.Errorf("node@%v exists but was not in node scan results", laddr)
diff --git a/lib/btrfsutil/old_rebuilt_forrest.go b/lib/btrfsutil/old_rebuilt_forrest.go
index 24faf79..36abe6f 100644
--- a/lib/btrfsutil/old_rebuilt_forrest.go
+++ b/lib/btrfsutil/old_rebuilt_forrest.go
@@ -231,7 +231,7 @@ func (tree oldRebuiltTree) addErrs(fn func(btrfsprim.Key, uint32) int, err error
 }
 
 func (bt *OldRebuiltForrest) readNode(nodeInfo nodeInfo) *btrfstree.Node {
-	node, err := bt.inner.ReadNode(bt.ctx, nodeInfo.LAddr, btrfstree.NodeExpectations{
+	node, err := bt.inner.AcquireNode(bt.ctx, nodeInfo.LAddr, btrfstree.NodeExpectations{
 		LAddr:      containers.OptionalValue(nodeInfo.LAddr),
 		Level:      containers.OptionalValue(nodeInfo.Level),
 		Generation: containers.OptionalValue(nodeInfo.Generation),
@@ -281,7 +281,7 @@ func (tree oldRebuiltTree) treeSearch(_ context.Context, searcher btrfstree.Tree
 	}
 
 	node := tree.forrest.readNode(indexItem.Value.Node)
-	defer node.Free()
+	defer tree.forrest.inner.ReleaseNode(node)
 
 	item := node.BodyLeaf[indexItem.Value.Slot]
 	item.Body = item.Body.CloneItem()
@@ -318,12 +318,12 @@ func (tree oldRebuiltTree) treeSubrange(_ context.Context, min int, searcher btr
 		func(rbNode *containers.RBNode[oldRebuiltTreeValue]) bool {
 			cnt++
 			if node == nil || node.Head.Addr != rbNode.Value.Node.LAddr {
-				node.Free()
+				tree.forrest.inner.ReleaseNode(node)
 				node = tree.forrest.readNode(rbNode.Value.Node)
 			}
 			return handleFn(node.BodyLeaf[rbNode.Value.Slot])
 		})
-	node.Free()
+	tree.forrest.inner.ReleaseNode(node)
 
 	var err error
 	if cnt < min {
@@ -366,7 +366,7 @@ func (tree oldRebuiltTree) treeWalk(ctx context.Context, cbs btrfstree.TreeWalkH
 			return false
 		}
 		if node == nil || node.Head.Addr != indexItem.Value.Node.LAddr {
-			node.Free()
+			tree.forrest.inner.ReleaseNode(node)
 			node = tree.forrest.readNode(indexItem.Value.Node)
 		}
 		item := node.BodyLeaf[indexItem.Value.Slot]
@@ -399,7 +399,7 @@ func (tree oldRebuiltTree) treeWalk(ctx context.Context, cbs btrfstree.TreeWalkH
 		}
 		return ctx.Err() == nil
 	})
-	node.Free()
+	tree.forrest.inner.ReleaseNode(node)
 }
 
 // Superblock implements btrfs.ReadableFS.
diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index 4f7a24b..60f3010 100644
--- a/lib/btrfsutil/rebuilt_forrest.go
+++ b/lib/btrfsutil/rebuilt_forrest.go
@@ -141,8 +141,6 @@ type RebuiltForrest struct {
 	leafs    containers.Cache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]
 	incItems containers.Cache[btrfsprim.ObjID, itemIndex]
 	excItems containers.Cache[btrfsprim.ObjID, itemIndex]
-
-	nodes containers.Cache[btrfsvol.LogicalAddr, btrfstree.Node]
 }
 
 // NewRebuiltForrest returns a new RebuiltForrest instance.  The
@@ -170,9 +168,6 @@ func NewRebuiltForrest(file btrfstree.NodeSource, sb btrfstree.Superblock, graph
 		containers.SourceFunc[btrfsprim.ObjID, itemIndex](func(ctx context.Context, treeID btrfsprim.ObjID, excItems *itemIndex) {
 			*excItems = ret.trees[treeID].uncachedExcItems(ctx)
 		}))
-	ret.nodes = containers.NewARCache[btrfsvol.LogicalAddr, btrfstree.Node](textui.Tunable(8),
-		containers.SourceFunc[btrfsvol.LogicalAddr, btrfstree.Node](ret.readNode))
-
 	if ret.cb == nil {
 		ret.cb = noopRebuiltForrestCallbacks{
 			forrest: ret,
diff --git a/lib/btrfsutil/rebuilt_readitem.go b/lib/btrfsutil/rebuilt_readitem.go
index 2520593..73dc01e 100644
--- a/lib/btrfsutil/rebuilt_readitem.go
+++ b/lib/btrfsutil/rebuilt_readitem.go
@@ -8,8 +8,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/datawire/dlib/dlog"
-
 	"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"
@@ -26,16 +24,20 @@ func (ptr ItemPtr) String() string {
 	return fmt.Sprintf("node@%v[%v]", ptr.Node, ptr.Slot)
 }
 
-func (ts *RebuiltForrest) readNode(ctx context.Context, laddr btrfsvol.LogicalAddr, out *btrfstree.Node) {
-	dlog.Debugf(ctx, "cache-miss node@%v, reading...", laddr)
-
-	graphInfo, ok := ts.graph.Nodes[laddr]
+func (ts *RebuiltForrest) readItem(ctx context.Context, ptr ItemPtr) btrfsitem.Item {
+	graphInfo, ok := ts.graph.Nodes[ptr.Node]
 	if !ok {
-		panic(fmt.Errorf("should not happen: node@%v is not mentioned in the in-memory graph", laddr))
+		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for node@%v not mentioned in the in-memory graph", ptr.Node))
+	}
+	if graphInfo.Level != 0 {
+		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for non-leaf node@%v", ptr.Node))
+	}
+	if ptr.Slot < 0 {
+		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for negative item slot: %v", ptr.Slot))
 	}
 
-	node, err := ts.file.ReadNode(ctx, laddr, btrfstree.NodeExpectations{
-		LAddr:      containers.OptionalValue(laddr),
+	node, err := ts.file.AcquireNode(ctx, ptr.Node, btrfstree.NodeExpectations{
+		LAddr:      containers.OptionalValue(ptr.Node),
 		Level:      containers.OptionalValue(graphInfo.Level),
 		Generation: containers.OptionalValue(graphInfo.Generation),
 		Owner: func(treeID btrfsprim.ObjID, gen btrfsprim.Generation) error {
@@ -49,22 +51,12 @@ func (ts *RebuiltForrest) readNode(ctx context.Context, laddr btrfsvol.LogicalAd
 		MinItem: containers.OptionalValue(graphInfo.MinItem()),
 		MaxItem: containers.OptionalValue(graphInfo.MaxItem()),
 	})
+	defer ts.file.ReleaseNode(node)
 	if err != nil {
 		panic(fmt.Errorf("should not happen: i/o error: %w", err))
 	}
-	*out = *node
-}
-
-func (ts *RebuiltForrest) readItem(ctx context.Context, ptr ItemPtr) btrfsitem.Item {
-	if ts.graph.Nodes[ptr.Node].Level != 0 {
-		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for non-leaf node@%v", ptr.Node))
-	}
-	if ptr.Slot < 0 {
-		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for negative item slot: %v", ptr.Slot))
-	}
 
-	items := ts.nodes.Acquire(ctx, ptr.Node).BodyLeaf
-	defer ts.nodes.Release(ptr.Node)
+	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",
diff --git a/lib/btrfsutil/scan.go b/lib/btrfsutil/scan.go
index 12e1dfe..0e268e5 100644
--- a/lib/btrfsutil/scan.go
+++ b/lib/btrfsutil/scan.go
@@ -134,7 +134,7 @@ func ScanOneDevice[Stats comparable, Result any](ctx context.Context, dev *btrfs
 				}
 				minNextNode = pos + btrfsvol.PhysicalAddr(sb.NodeSize)
 			}
-			node.Free()
+			node.RawFree()
 		}
 	}