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

1 files changed, 46 insertions, 11 deletions

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)