Merge branch 'lukeshu/node-cache'

8 files changed, 166 insertions, 88 deletions

diff --git a/lib/btrfs/btrfstree/btree.go b/lib/btrfs/btrfstree/btree.go
index bc1bac2..4f5d21b 100644
--- a/lib/btrfs/btrfstree/btree.go
+++ b/lib/btrfs/btrfstree/btree.go
@@ -11,6 +11,7 @@ import (
 	"fmt"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
 type TreeSearcher interface {
@@ -96,5 +97,6 @@ func (e *TreeError) Error() string {
 
 type NodeSource interface {
 	Superblock() (*Superblock, error)
-	ReadNode(Path) (*Node, error)
+	AcquireNode(ctx context.Context, addr btrfsvol.LogicalAddr, exp NodeExpectations) (*Node, error)
+	ReleaseNode(*Node)
 }
diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index 38a2721..b04bfc0 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -83,11 +83,14 @@ func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID bt
 }
 
 type TreeOperatorImpl struct {
-	NodeSource
+	NodeSource interface {
+		NodeSource
+		NodeFile
+	}
 }
 
 func (fs TreeOperatorImpl) RawTree(ctx context.Context, treeID btrfsprim.ObjID) (*RawTree, error) {
-	sb, err := fs.Superblock()
+	sb, err := fs.NodeSource.Superblock()
 	if err != nil {
 		return nil, err
 	}
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index a943803..86eab11 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -22,6 +22,9 @@ type RawTree struct {
 }
 
 func (tree *RawTree) TreeWalk(ctx context.Context, cbs TreeWalkHandler) {
+	if tree.RootNode == 0 {
+		return
+	}
 	path := Path{{
 		FromTree:         tree.ID,
 		FromItemSlot:     -1,
@@ -37,13 +40,14 @@ func (tree *RawTree) walk(ctx context.Context, path Path, cbs TreeWalkHandler) {
 	if ctx.Err() != nil {
 		return
 	}
-	if path.Node(-1).ToNodeAddr == 0 {
-		return
-	}
 
 	// 001
-	node, err := tree.Forrest.ReadNode(path)
-	defer node.Free()
+	nodeAddr, nodeExp, ok := path.NodeExpectations(tree.Forrest.NodeSource)
+	if !ok {
+		return
+	}
+	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/node_exp.go b/lib/btrfs/btrfstree/node_exp.go
new file mode 100644
index 0000000..dec6dbb
--- /dev/null
+++ b/lib/btrfs/btrfstree/node_exp.go
@@ -0,0 +1,67 @@
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfstree
+
+import (
+	"fmt"
+
+	"github.com/datawire/dlib/derror"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+)
+
+type NodeExpectations struct {
+	LAddr containers.Optional[btrfsvol.LogicalAddr]
+	// Things knowable from the parent.
+	Level      containers.Optional[uint8]
+	Generation containers.Optional[btrfsprim.Generation]
+	Owner      func(btrfsprim.ObjID, btrfsprim.Generation) error
+	MinItem    containers.Optional[btrfsprim.Key]
+	// Things knowable from the structure of the tree.
+	MaxItem containers.Optional[btrfsprim.Key]
+}
+
+func (exp NodeExpectations) Check(node *Node) error {
+	var errs derror.MultiError
+	if exp.LAddr.OK && node.Head.Addr != exp.LAddr.Val {
+		errs = append(errs, fmt.Errorf("read from laddr=%v but claims to be at laddr=%v",
+			exp.LAddr.Val, node.Head.Addr))
+	}
+	if exp.Level.OK && node.Head.Level != exp.Level.Val {
+		errs = append(errs, fmt.Errorf("expected level=%v but claims to be level=%v",
+			exp.Level.Val, node.Head.Level))
+	}
+	if node.Head.Level > MaxLevel {
+		errs = append(errs, fmt.Errorf("maximum level=%v but claims to be level=%v",
+			MaxLevel, node.Head.Level))
+	}
+	if exp.Generation.OK && node.Head.Generation != exp.Generation.Val {
+		errs = append(errs, fmt.Errorf("expected generation=%v but claims to be generation=%v",
+			exp.Generation.Val, node.Head.Generation))
+	}
+	if exp.Owner != nil {
+		if err := exp.Owner(node.Head.Owner, node.Head.Generation); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	if node.Head.NumItems == 0 {
+		errs = append(errs, fmt.Errorf("has no items"))
+	} else {
+		if minItem, _ := node.MinItem(); exp.MinItem.OK && exp.MinItem.Val.Compare(minItem) > 0 {
+			errs = append(errs, fmt.Errorf("expected minItem>=%v but node has minItem=%v",
+				exp.MinItem, minItem))
+		}
+		if maxItem, _ := node.MaxItem(); exp.MaxItem.OK && exp.MaxItem.Val.Compare(maxItem) < 0 {
+			errs = append(errs, fmt.Errorf("expected maxItem<=%v but node has maxItem=%v",
+				exp.MaxItem, maxItem))
+		}
+	}
+	if len(errs) > 0 {
+		return errs
+	}
+	return nil
+}
diff --git a/lib/btrfs/btrfstree/readnode.go b/lib/btrfs/btrfstree/readnode.go
index ac82c62..a4ccf10 100644
--- a/lib/btrfs/btrfstree/readnode.go
+++ b/lib/btrfs/btrfstree/readnode.go
@@ -14,8 +14,7 @@ import (
 )
 
 type NodeFile interface {
-	diskio.File[btrfsvol.LogicalAddr]
-	Superblock() (*Superblock, error)
+	diskio.ReaderAt[btrfsvol.LogicalAddr]
 
 	// ParentTree, given a tree ID, returns that tree's parent
 	// tree, if it has one.
@@ -29,15 +28,14 @@ type NodeFile interface {
 	ParentTree(btrfsprim.ObjID) (btrfsprim.ObjID, btrfsprim.Generation, bool)
 }
 
-// FSReadNode is a utility function to help with implementing the
-// 'NodeSource' interface.
-func FSReadNode(
-	fs NodeFile,
-	path Path,
-) (*Node, error) {
-	sb, err := fs.Superblock()
-	if err != nil {
-		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
+// NodeExpectations returns the address to read and the expectations
+// to have when reading the node pointed to by this Path.
+//
+// `ok` is false if the path is empty or if this Path points to an
+// item rather than a node.
+func (path Path) NodeExpectations(fs NodeFile) (_ btrfsvol.LogicalAddr, _ NodeExpectations, ok bool) {
+	if path.Node(-1).ToNodeAddr == 0 && path.Node(-1).ToNodeGeneration == 0 && path.Node(-1).ToNodeLevel == 0 {
+		return 0, NodeExpectations{}, false
 	}
 
 	checkOwner := func(owner btrfsprim.ObjID, gen btrfsprim.Generation) error {
@@ -70,12 +68,12 @@ func FSReadNode(
 		}
 	}
 
-	return ReadNode[btrfsvol.LogicalAddr](fs, *sb, path.Node(-1).ToNodeAddr, NodeExpectations{
+	return path.Node(-1).ToNodeAddr, NodeExpectations{
 		LAddr:      containers.OptionalValue(path.Node(-1).ToNodeAddr),
 		Level:      containers.OptionalValue(path.Node(-1).ToNodeLevel),
 		Generation: containers.OptionalValue(path.Node(-1).ToNodeGeneration),
 		Owner:      checkOwner,
 		MinItem:    containers.OptionalValue(path.Node(-1).ToKey),
 		MaxItem:    containers.OptionalValue(path.Node(-1).ToMaxKey),
-	})
+	}, true
 }
diff --git a/lib/btrfs/btrfstree/types_node.go b/lib/btrfs/btrfstree/types_node.go
index bfcbbf4..e29431f 100644
--- a/lib/btrfs/btrfstree/types_node.go
+++ b/lib/btrfs/btrfstree/types_node.go
@@ -10,7 +10,6 @@ import (
 	"fmt"
 
 	"git.lukeshu.com/go/typedsync"
-	"github.com/datawire/dlib/derror"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
@@ -104,10 +103,13 @@ type NodeHeader struct {
 	Generation    btrfsprim.Generation `bin:"off=0x50, siz=0x8"`
 	Owner         btrfsprim.ObjID      `bin:"off=0x58, siz=0x8"` // The ID of the tree that contains this node
 	NumItems      uint32               `bin:"off=0x60, siz=0x4"` // [ignored-when-writing]
-	Level         uint8                `bin:"off=0x64, siz=0x1"` // 0 for leaf nodes, >=1 for interior nodes
+	Level         uint8                `bin:"off=0x64, siz=0x1"` // 0 for leaf nodes, >=1 for interior nodes (max 8)
 	binstruct.End `bin:"off=0x65"`
 }
 
+// MaxLevel is the maximum valid NodeHeader.Level.
+const MaxLevel = 8
+
 // MaxItems returns the maximum possible valid value of
 // .Head.NumItems.
 func (node Node) MaxItems() uint32 {
@@ -306,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
 	}
@@ -412,17 +416,6 @@ func (node *Node) LeafFreeSpace() uint32 {
 
 var ErrNotANode = errors.New("does not look like a node")
 
-type NodeExpectations struct {
-	LAddr containers.Optional[btrfsvol.LogicalAddr]
-	// Things knowable from the parent.
-	Level      containers.Optional[uint8]
-	Generation containers.Optional[btrfsprim.Generation]
-	Owner      func(btrfsprim.ObjID, btrfsprim.Generation) error
-	MinItem    containers.Optional[btrfsprim.Key]
-	// Things knowable from the structure of the tree.
-	MaxItem containers.Optional[btrfsprim.Key]
-}
-
 type NodeError[Addr ~int64] struct {
 	Op       string
 	NodeAddr Addr
@@ -455,7 +448,7 @@ var nodePool = typedsync.Pool[*Node]{
 // returned.  The error returned (if non-nil) is always of type
 // *NodeError[Addr].  Notable errors that may be inside of the
 // NodeError are ErrNotANode and *IOError.
-func ReadNode[Addr ~int64](fs diskio.ReaderAt[Addr], sb Superblock, addr Addr, exp NodeExpectations) (*Node, error) {
+func ReadNode[Addr ~int64](fs diskio.ReaderAt[Addr], sb Superblock, addr Addr) (*Node, error) {
 	if int(sb.NodeSize) < nodeHeaderSize {
 		return nil, &NodeError[Addr]{
 			Op: "btrfstree.ReadNode", NodeAddr: addr,
@@ -521,50 +514,7 @@ func ReadNode[Addr ~int64](fs diskio.ReaderAt[Addr], sb Superblock, addr Addr, e
 
 	bytePool.Put(nodeBuf)
 
-	// sanity checking (that doesn't prevent parsing)
-
-	if err := exp.Check(node); err != nil {
-		return node, &NodeError[Addr]{Op: "btrfstree.ReadNode", NodeAddr: addr, Err: err}
-	}
-
 	// return
 
 	return node, nil
 }
-
-func (exp NodeExpectations) Check(node *Node) error {
-	var errs derror.MultiError
-	if exp.LAddr.OK && node.Head.Addr != exp.LAddr.Val {
-		errs = append(errs, fmt.Errorf("read from laddr=%v but claims to be at laddr=%v",
-			exp.LAddr.Val, node.Head.Addr))
-	}
-	if exp.Level.OK && node.Head.Level != exp.Level.Val {
-		errs = append(errs, fmt.Errorf("expected level=%v but claims to be level=%v",
-			exp.Level.Val, node.Head.Level))
-	}
-	if exp.Generation.OK && node.Head.Generation != exp.Generation.Val {
-		errs = append(errs, fmt.Errorf("expected generation=%v but claims to be generation=%v",
-			exp.Generation.Val, node.Head.Generation))
-	}
-	if exp.Owner != nil {
-		if err := exp.Owner(node.Head.Owner, node.Head.Generation); err != nil {
-			errs = append(errs, err)
-		}
-	}
-	if node.Head.NumItems == 0 {
-		errs = append(errs, fmt.Errorf("has no items"))
-	} else {
-		if minItem, _ := node.MinItem(); exp.MinItem.OK && exp.MinItem.Val.Compare(minItem) > 0 {
-			errs = append(errs, fmt.Errorf("expected minItem>=%v but node has minItem=%v",
-				exp.MinItem, minItem))
-		}
-		if maxItem, _ := node.MaxItem(); exp.MaxItem.OK && exp.MaxItem.Val.Compare(maxItem) < 0 {
-			errs = append(errs, fmt.Errorf("expected maxItem<=%v but node has maxItem=%v",
-				exp.MaxItem, maxItem))
-		}
-	}
-	if len(errs) > 0 {
-		return errs
-	}
-	return nil
-}
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 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.