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

5 files changed, 56 insertions, 15 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)