1 files changed, 60 insertions, 9 deletions
diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index b60f54a..8aa485f 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -6,23 +6,18 @@ package btrfs
 
 import (
 	"context"
+	"fmt"
 
 	"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"
+	"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.
 
-// btrfstree.NodeSource ////////////////////////////////////////////////////////
-
-// ReadNode implements btrfstree.NodeSource.
-func (fs *FS) ReadNode(path btrfstree.Path) (*btrfstree.Node, error) {
-	return btrfstree.FSReadNode(fs, path)
-}
-
-var _ btrfstree.NodeSource = (*FS)(nil)
-
 // btrfstree.NodeFile //////////////////////////////////////////////////////////
 
 type treeInfo struct {
@@ -85,6 +80,62 @@ func (fs *FS) ParentTree(tree btrfsprim.ObjID) (btrfsprim.ObjID, btrfsprim.Gener
 
 var _ btrfstree.NodeFile = (*FS)(nil)
 
+// btrfstree.NodeSource ////////////////////////////////////////////////////////
+
+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),
+		)
+	}
+
+	nodeEntry := fs.cacheNodes.Acquire(ctx, addr)
+	if nodeEntry.err != nil {
+		err := nodeEntry.err
+		fs.cacheNodes.Release(addr)
+		return nil, 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
+	}
+
+	nodeEntry.node, nodeEntry.err = btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, addr)
+}
+
+var _ btrfstree.NodeSource = (*FS)(nil)
+
 // btrfstree.TreeOperator //////////////////////////////////////////////////////
 
 // TreeWalk implements btrfstree.TreeOperator.