btrfstree: Add a FreeNodeRef function, use it

7 files changed, 99 insertions, 9 deletions

diff --git a/lib/btrfs/btrfstree/ops.go b/lib/btrfs/btrfstree/ops.go
index bda4ac8..b01312f 100644
--- a/lib/btrfs/btrfstree/ops.go
+++ b/lib/btrfs/btrfstree/ops.go
@@ -144,6 +144,7 @@ func (fs TreeOperatorImpl) treeWalk(ctx context.Context, path TreePath, errHandl
 		}
 	}
 	node, err := fs.ReadNode(path)
+	defer FreeNodeRef(node)
 	if ctx.Err() != nil {
 		return
 	}
@@ -255,6 +256,7 @@ func (fs TreeOperatorImpl) treeSearch(treeRoot TreeRoot, fn func(btrfsprim.Key,
 		}
 		node, err := fs.ReadNode(path)
 		if err != nil {
+			FreeNodeRef(node)
 			return nil, nil, err
 		}
 
@@ -273,6 +275,7 @@ func (fs TreeOperatorImpl) treeSearch(treeRoot TreeRoot, fn func(btrfsprim.Key,
 				return slices.Min(fn(kp.Key, math.MaxUint32), 0) // don't return >0; a key can't be "too low"
 			})
 			if !ok {
+				FreeNodeRef(node)
 				return nil, nil, iofs.ErrNotExist
 			}
 			toMaxKey := path.Node(-1).ToMaxKey
@@ -288,6 +291,7 @@ func (fs TreeOperatorImpl) treeSearch(treeRoot TreeRoot, fn func(btrfsprim.Key,
 				ToKey:            node.Data.BodyInternal[lastGood].Key,
 				ToMaxKey:         toMaxKey,
 			})
+			FreeNodeRef(node)
 		} else {
 			// leaf node
 
@@ -305,6 +309,7 @@ func (fs TreeOperatorImpl) treeSearch(treeRoot TreeRoot, fn func(btrfsprim.Key,
 				return fn(item.Key, item.BodySize)
 			})
 			if !ok {
+				FreeNodeRef(node)
 				return nil, nil, iofs.ErrNotExist
 			}
 			path = append(path, TreePathElem{
@@ -333,8 +338,10 @@ func (fs TreeOperatorImpl) prev(path TreePath, node *diskio.Ref[btrfsvol.Logical
 	path.Node(-1).FromItemIdx--
 	if path.Node(-1).ToNodeAddr != 0 {
 		if node.Addr != path.Node(-2).ToNodeAddr {
+			FreeNodeRef(node)
 			node, err = fs.ReadNode(path.Parent())
 			if err != nil {
+				FreeNodeRef(node)
 				return nil, nil, err
 			}
 			path.Node(-1).ToNodeAddr = node.Data.BodyInternal[path.Node(-1).FromItemIdx].BlockPtr
@@ -343,8 +350,10 @@ func (fs TreeOperatorImpl) prev(path TreePath, node *diskio.Ref[btrfsvol.Logical
 	// go down
 	for path.Node(-1).ToNodeAddr != 0 {
 		if node.Addr != path.Node(-1).ToNodeAddr {
+			FreeNodeRef(node)
 			node, err = fs.ReadNode(path)
 			if err != nil {
+				FreeNodeRef(node)
 				return nil, nil, err
 			}
 		}
@@ -369,8 +378,10 @@ func (fs TreeOperatorImpl) prev(path TreePath, node *diskio.Ref[btrfsvol.Logical
 	}
 	// return
 	if node.Addr != path.Node(-2).ToNodeAddr {
+		FreeNodeRef(node)
 		node, err = fs.ReadNode(path.Parent())
 		if err != nil {
+			FreeNodeRef(node)
 			return nil, nil, err
 		}
 	}
@@ -383,8 +394,10 @@ func (fs TreeOperatorImpl) next(path TreePath, node *diskio.Ref[btrfsvol.Logical
 
 	// go up
 	if node.Addr != path.Node(-2).ToNodeAddr {
+		FreeNodeRef(node)
 		node, err = fs.ReadNode(path.Parent())
 		if err != nil {
+			FreeNodeRef(node)
 			return nil, nil, err
 		}
 		path.Node(-2).ToNodeLevel = node.Data.Head.Level
@@ -395,8 +408,10 @@ func (fs TreeOperatorImpl) next(path TreePath, node *diskio.Ref[btrfsvol.Logical
 			return nil, nil, nil
 		}
 		if node.Addr != path.Node(-2).ToNodeAddr {
+			FreeNodeRef(node)
 			node, err = fs.ReadNode(path.Parent())
 			if err != nil {
+				FreeNodeRef(node)
 				return nil, nil, err
 			}
 			path.Node(-2).ToNodeLevel = node.Data.Head.Level
@@ -406,8 +421,10 @@ func (fs TreeOperatorImpl) next(path TreePath, node *diskio.Ref[btrfsvol.Logical
 	path.Node(-1).FromItemIdx++
 	if path.Node(-1).ToNodeAddr != 0 {
 		if node.Addr != path.Node(-2).ToNodeAddr {
+			FreeNodeRef(node)
 			node, err = fs.ReadNode(path.Parent())
 			if err != nil {
+				FreeNodeRef(node)
 				return nil, nil, err
 			}
 			path.Node(-1).ToNodeAddr = node.Data.BodyInternal[path.Node(-1).FromItemIdx].BlockPtr
@@ -416,8 +433,10 @@ func (fs TreeOperatorImpl) next(path TreePath, node *diskio.Ref[btrfsvol.Logical
 	// go down
 	for path.Node(-1).ToNodeAddr != 0 {
 		if node.Addr != path.Node(-1).ToNodeAddr {
+			FreeNodeRef(node)
 			node, err = fs.ReadNode(path)
 			if err != nil {
+				FreeNodeRef(node)
 				return nil, nil, err
 			}
 			path.Node(-1).ToNodeLevel = node.Data.Head.Level
@@ -447,8 +466,10 @@ func (fs TreeOperatorImpl) next(path TreePath, node *diskio.Ref[btrfsvol.Logical
 	}
 	// return
 	if node.Addr != path.Node(-2).ToNodeAddr {
+		FreeNodeRef(node)
 		node, err = fs.ReadNode(path.Parent())
 		if err != nil {
+			FreeNodeRef(node)
 			return nil, nil, err
 		}
 	}
@@ -469,7 +490,10 @@ func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, fn func(btrfsprim.
 	if err != nil {
 		return Item{}, err
 	}
-	return node.Data.BodyLeaf[path.Node(-1).FromItemIdx], nil
+	item := node.Data.BodyLeaf[path.Node(-1).FromItemIdx]
+	item.Body = item.Body.CloneItem()
+	FreeNodeRef(node)
+	return item, nil
 }
 
 // KeySearch returns a comparator suitable to be passed to TreeSearch.
@@ -506,7 +530,8 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 
 	ret := []Item{middleItem}
 	var errs derror.MultiError
-	for prevPath, prevNode := middlePath, middleNode; true; {
+	prevPath, prevNode := middlePath, middleNode
+	for {
 		prevPath, prevNode, err = fs.prev(prevPath, prevNode)
 		if err != nil {
 			errs = append(errs, err)
@@ -519,10 +544,21 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 		if fn(prevItem.Key, prevItem.BodySize) != 0 {
 			break
 		}
-		ret = append(ret, prevItem)
+		item := prevItem
+		item.Body = item.Body.CloneItem()
+		ret = append(ret, item)
 	}
 	slices.Reverse(ret)
-	for nextPath, nextNode := middlePath, middleNode; true; {
+	if prevNode.Addr != middlePath.Node(-1).ToNodeAddr {
+		FreeNodeRef(prevNode)
+		middleNode, err = fs.ReadNode(middlePath)
+		if err != nil {
+			FreeNodeRef(middleNode)
+			return nil, err
+		}
+	}
+	nextPath, nextNode := middlePath, middleNode
+	for {
 		nextPath, nextNode, err = fs.next(nextPath, nextNode)
 		if err != nil {
 			errs = append(errs, err)
@@ -535,8 +571,11 @@ func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfspr
 		if fn(nextItem.Key, nextItem.BodySize) != 0 {
 			break
 		}
-		ret = append(ret, nextItem)
+		item := nextItem
+		item.Body = item.Body.CloneItem()
+		ret = append(ret, item)
 	}
+	FreeNodeRef(nextNode)
 	if errs != nil {
 		err = errs
 	}
diff --git a/lib/btrfs/btrfstree/types_node.go b/lib/btrfs/btrfstree/types_node.go
index 85aae23..fd4c939 100644
--- a/lib/btrfs/btrfstree/types_node.go
+++ b/lib/btrfs/btrfstree/types_node.go
@@ -8,7 +8,9 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"unsafe"
 
+	"git.lukeshu.com/go/typedsync"
 	"github.com/datawire/dlib/derror"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct"
@@ -300,12 +302,24 @@ type ItemHeader struct {
 	binstruct.End `bin:"off=0x19"`
 }
 
+var itemPool containers.SlicePool[Item]
+
+func (node *Node) Free() {
+	for i := range node.BodyLeaf {
+		node.BodyLeaf[i].Body.Free()
+		node.BodyLeaf[i] = Item{}
+	}
+	itemPool.Put(node.BodyLeaf)
+	*node = Node{}
+}
+
 func (node *Node) unmarshalLeaf(bodyBuf []byte) (int, error) {
 	head := 0
 	tail := len(bodyBuf)
-	node.BodyLeaf = make([]Item, node.Head.NumItems)
+	node.BodyLeaf = itemPool.Get(int(node.Head.NumItems))
+	var itemHead ItemHeader
 	for i := range node.BodyLeaf {
-		var itemHead ItemHeader
+		itemHead = ItemHeader{} // zero it out
 		n, err := binstruct.Unmarshal(bodyBuf[head:], &itemHead)
 		head += n
 		if err != nil {
@@ -423,6 +437,25 @@ func (e *IOError) Unwrap() error { return e.Err }
 
 var bytePool containers.SlicePool[byte]
 
+var nodePool = typedsync.Pool[*diskio.Ref[int64, Node]]{
+	New: func() *diskio.Ref[int64, Node] {
+		return new(diskio.Ref[int64, Node])
+	},
+}
+
+func FreeNodeRef[Addr ~int64](ref *diskio.Ref[Addr, Node]) {
+	if ref == nil {
+		return
+	}
+	ref.Data.Free()
+	nodePool.Put((*diskio.Ref[int64, Node])(unsafe.Pointer(ref))) //nolint:gosec // I know it's unsafe.
+}
+
+func newNodeRef[Addr ~int64]() *diskio.Ref[Addr, Node] {
+	ret, _ := nodePool.Get()
+	return (*diskio.Ref[Addr, Node])(unsafe.Pointer(ret)) //nolint:gosec // I know it's unsafe.
+}
+
 // It is possible that both a non-nil diskio.Ref and an error are
 // returned.  The error returned (if non-nil) is always of type
 // *NodeError[Addr].  Notable errors that may be inside of the
@@ -443,7 +476,8 @@ func ReadNode[Addr ~int64](fs diskio.File[Addr], sb Superblock, addr Addr, exp N
 
 	// parse (early)
 
-	nodeRef := &diskio.Ref[Addr, Node]{
+	nodeRef := newNodeRef[Addr]()
+	*nodeRef = diskio.Ref[Addr, Node]{
 		File: fs,
 		Addr: addr,
 		Data: Node{
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index 149706d..b4ab645 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -62,6 +62,9 @@ func NewHandle(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock)
 
 		cache: containers.ARCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]{
 			MaxLen: textui.Tunable(8),
+			OnRemove: func(_ btrfsvol.LogicalAddr, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
+				btrfstree.FreeNodeRef(nodeRef)
+			},
 		},
 	}
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 7e19802..632ed70 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -1,4 +1,4 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
@@ -48,12 +48,15 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			})
 			if err != nil {
+				btrfstree.FreeNodeRef(nodeRef)
 				return btrfstree.Superblock{}, graph.Graph{}, nil, err
 			}
 
 			nodeGraph.InsertNode(nodeRef)
 			keyIO.InsertNode(nodeRef)
 
+			btrfstree.FreeNodeRef(nodeRef)
+
 			stats.N++
 			progressWriter.Set(stats)
 		}
diff --git a/lib/btrfsprogs/btrfsinspect/scandevices.go b/lib/btrfsprogs/btrfsinspect/scandevices.go
index 0cfee5b..d54be71 100644
--- a/lib/btrfsprogs/btrfsinspect/scandevices.go
+++ b/lib/btrfsprogs/btrfsinspect/scandevices.go
@@ -245,6 +245,7 @@ func ScanOneDevice(ctx context.Context, dev *btrfs.Device, sb btrfstree.Superblo
 				}
 				minNextNode = pos + btrfsvol.PhysicalAddr(sb.NodeSize)
 			}
+			btrfstree.FreeNodeRef(nodeRef)
 		}
 	}
 	progress(devSize)
diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go
index 7ea31ce..15641ab 100644
--- a/lib/btrfsprogs/btrfsutil/broken_btree.go
+++ b/lib/btrfsprogs/btrfsutil/broken_btree.go
@@ -237,11 +237,13 @@ func (bt *brokenTrees) TreeSearch(treeID btrfsprim.ObjID, fn func(btrfsprim.Key,
 
 	itemPath := bt.arena.Inflate(indexItem.Value.Path)
 	node, err := bt.inner.ReadNode(itemPath.Parent())
+	defer btrfstree.FreeNodeRef(node)
 	if err != nil {
 		return btrfstree.Item{}, bt.addErrs(index, fn, err)
 	}
 
 	item := node.Data.BodyLeaf[itemPath.Node(-1).FromItemIdx]
+	item.Body = item.Body.CloneItem()
 
 	// Since we were only asked to return 1 item, it isn't
 	// necessary to augment this `nil` with bt.addErrs().
@@ -271,13 +273,16 @@ func (bt *brokenTrees) TreeSearchAll(treeID btrfsprim.ObjID, fn func(btrfsprim.K
 		itemPath := bt.arena.Inflate(indexItems[i].Path)
 		if node == nil || node.Addr != itemPath.Node(-2).ToNodeAddr {
 			var err error
+			btrfstree.FreeNodeRef(node)
 			node, err = bt.inner.ReadNode(itemPath.Parent())
 			if err != nil {
+				btrfstree.FreeNodeRef(node)
 				return nil, bt.addErrs(index, fn, err)
 			}
 		}
 		ret[i] = node.Data.BodyLeaf[itemPath.Node(-1).FromItemIdx]
 	}
+	btrfstree.FreeNodeRef(node)
 
 	return ret, bt.addErrs(index, fn, nil)
 }
@@ -306,8 +311,10 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, err
 			itemPath := bt.arena.Inflate(indexItem.Value.Path)
 			if node == nil || node.Addr != itemPath.Node(-2).ToNodeAddr {
 				var err error
+				btrfstree.FreeNodeRef(node)
 				node, err = bt.inner.ReadNode(itemPath.Parent())
 				if err != nil {
+					btrfstree.FreeNodeRef(node)
 					errHandle(&btrfstree.TreeError{Path: itemPath, Err: err})
 					return true
 				}
@@ -319,6 +326,7 @@ func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, err
 		}
 		return true
 	})
+	btrfstree.FreeNodeRef(node)
 }
 
 func (bt *brokenTrees) Superblock() (*btrfstree.Superblock, error) {
@@ -339,6 +347,7 @@ func (bt *brokenTrees) Augment(treeID btrfsprim.ObjID, nodeAddr btrfsvol.Logical
 		return nil, index.TreeRootErr
 	}
 	nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](bt.inner, *sb, nodeAddr, btrfstree.NodeExpectations{})
+	defer btrfstree.FreeNodeRef(nodeRef)
 	if err != nil {
 		return nil, err
 	}
diff --git a/lib/btrfsprogs/btrfsutil/skinny_paths.go b/lib/btrfsprogs/btrfsutil/skinny_paths.go
index 4c314ec..1695990 100644
--- a/lib/btrfsprogs/btrfsutil/skinny_paths.go
+++ b/lib/btrfsprogs/btrfsutil/skinny_paths.go
@@ -55,6 +55,7 @@ func (a *SkinnyPathArena) getItem(parent btrfstree.TreePath, itemIdx int) (btrfs
 	}
 
 	node, err := btrfstree.ReadNode(a.FS, a.SB, parent.Node(-1).ToNodeAddr, btrfstree.NodeExpectations{})
+	defer btrfstree.FreeNodeRef(node)
 	if err != nil {
 		return btrfstree.TreePathElem{}, err
 	}