btrfs: Rethink the ReadNode API to better encourage sanity checking

2 files changed, 105 insertions, 74 deletions

diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index 7404b6c..682e355 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -16,6 +16,7 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"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"
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
@@ -73,45 +74,48 @@ var _ Trees = (*FS)(nil)
 //
 //	[superblock]
 //	     |
-//	     | <------------------------------------------ pathElem={idx:-1, addr:0x01, lvl:3}
+//	     | <------------------------------------------ pathElem={idx:-1, addr:0x01, lvl:3, gen:6}
 //	     |
 //	  +[0x01]-----------+
-//	  | lvl=3           |
+//	  | lvl=3  gen=6    |
 //	  +-+-+-+-+-+-+-+-+-+
 //	  |1|2|3|4|5|6|7|8|9|
 //	  +---+---+---+---+-+
 //	                 |
-//	                 | <------------------------------ pathElem={idx:8, addr:0x02, lvl:2}
+//	                 | <------------------------------ pathElem={idx:8, addr:0x02, lvl:2: gen:6}
 //	                 |
 //	              +[0x02]-----------+
-//	              | lvl=2           |
+//	              | lvl=2  gen=5    |
 //	              +-+-+-+-+-+-+-+-+-+
 //	              |1|2|3|4|5|6|7|8|9|
 //	              +---+---+---+---+-+
 //	                           |
-//	                           | <-------------------- pathElem={idx:7, addr:0x03, lvl:1}
+//	                           | <-------------------- pathElem={idx:7, addr:0x03, lvl:1: gen:5}
 //	                           |
 //	                        +[0x03]-----------+
-//	                        | lvl=1           |
+//	                        | lvl=1  gen=5    |
 //	                        +-+-+-+-+-+-+-+-+-+
 //	                        |1|2|3|4|5|6|7|8|9|
 //	                        +---+---+---+---+-+
 //	                               |
-//	                               | <---------------- pathElem={idx:4, addr:0x04, lvl:0}
+//	                               | <---------------- pathElem={idx:4, addr:0x04, lvl:0: gen:5}
 //	                               |
 //	                             +[0x04]-----------+
-//	                             | lvl=0           |
+//	                             | lvl=0  gen=2    |
 //	                             +-+-+-+-+-+-+-+-+-+
 //	                             |1|2|3|4|5|6|7|8|9|
 //	                             +---+---+---+---+-+
 //	                                |
-//	                                | <--------------- pathElem={idx:5, addr:0, lvl:0}
+//	                                | <--------------- pathElem={idx:5, addr:0, lvl:0, gen:2}
 //	                                |
 //	                              [item]
 //
 // the path would be
 //
-//	{-1, 0x01, 3}→{8, 0x02, 2}→{7, 0x03, 1}→{4, 0x04, 0}→{2, 0, 0}
+//	{
+//		TreeID: …,
+//		Nodes: {-1, 0x01, 3, 6}→{8, 0x02, 2, 6}→{7, 0x03, 1, 5}→{4, 0x04, 5}→{2, 0, 0, 2},
+//	}
 type TreePath struct {
 	TreeID ObjID
 	Nodes  []TreePathElem
@@ -129,6 +133,9 @@ type TreePathElem struct {
 	// NodeLevel is the expected or actual level of the node at
 	// NodeAddr.
 	NodeLevel uint8
+	// NodeParentGeneration is the generation of the parent node;
+	// the maximum expected generation of the child node.
+	NodeParentGeneration Generation
 }
 
 func (elem TreePathElem) writeNodeTo(w io.Writer) {
@@ -165,6 +172,11 @@ func (path TreePath) Append(elem TreePathElem) TreePath {
 	return path
 }
 
+func (path TreePath) Parent() TreePath {
+	path.Nodes = path.Nodes[:len(path.Nodes)-1]
+	return path
+}
+
 // path.Node(x) is like path.Nodes[x], but negative values of x move
 // down from the end of path.Nodes (similar to how lists work in many
 // other languages, such as Python).
@@ -175,6 +187,20 @@ func (path TreePath) Node(x int) *TreePathElem {
 	return &path.Nodes[x]
 }
 
+func (fs *FS) ReadNode(path TreePath) (*diskio.Ref[btrfsvol.LogicalAddr, Node], error) {
+	sb, err := fs.Superblock()
+	if err != nil {
+		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
+	}
+
+	return ReadNode[btrfsvol.LogicalAddr](fs, *sb, path.Node(-1).NodeAddr, NodeExpectations{
+		LAddr:         containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: path.Node(-1).NodeAddr},
+		Level:         containers.Optional[uint8]{OK: true, Val: path.Node(-1).NodeLevel},
+		MaxGeneration: containers.Optional[Generation]{OK: true, Val: path.Node(-1).NodeParentGeneration},
+		Owner:         containers.Optional[ObjID]{OK: true, Val: path.TreeID},
+	})
+}
+
 type TreeError struct {
 	Path TreePath
 	Err  error
@@ -272,6 +298,7 @@ type TreeWalkHandler struct {
 	BadItem func(TreePath, Item) error
 }
 
+// TreeWalk implements the 'Trees' interface.
 func (fs *FS) TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
 	path := TreePath{
 		TreeID: treeID,
@@ -282,9 +309,10 @@ func (fs *FS) TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeEr
 		return
 	}
 	path = path.Append(TreePathElem{
-		ItemIdx:   -1,
-		NodeAddr:  rootInfo.RootNode,
-		NodeLevel: rootInfo.Level,
+		ItemIdx:              -1,
+		NodeAddr:             rootInfo.RootNode,
+		NodeLevel:            rootInfo.Level,
+		NodeParentGeneration: rootInfo.Generation,
 	})
 	fs.treeWalk(ctx, path, errHandle, cbs)
 }
@@ -296,9 +324,10 @@ func (fs *FS) RawTreeWalk(ctx context.Context, rootInfo TreeRoot, errHandle func
 		TreeID: rootInfo.TreeID,
 		Nodes: []TreePathElem{
 			{
-				ItemIdx:   -1,
-				NodeAddr:  rootInfo.RootNode,
-				NodeLevel: rootInfo.Level,
+				ItemIdx:              -1,
+				NodeAddr:             rootInfo.RootNode,
+				NodeLevel:            rootInfo.Level,
+				NodeParentGeneration: rootInfo.Generation,
 			},
 		},
 	}
@@ -321,7 +350,7 @@ func (fs *FS) treeWalk(ctx context.Context, path TreePath, errHandle func(*TreeE
 			return
 		}
 	}
-	node, err := fs.readNodeAtLevel(path.Node(-1).NodeAddr, path.Node(-1).NodeLevel)
+	node, err := fs.ReadNode(path)
 	if ctx.Err() != nil {
 		return
 	}
@@ -344,9 +373,10 @@ func (fs *FS) treeWalk(ctx context.Context, path TreePath, errHandle func(*TreeE
 	if node != nil {
 		for i, item := range node.Data.BodyInternal {
 			itemPath := path.Append(TreePathElem{
-				ItemIdx:   i,
-				NodeAddr:  item.BlockPtr,
-				NodeLevel: node.Data.Head.Level - 1,
+				ItemIdx:              i,
+				NodeAddr:             item.BlockPtr,
+				NodeLevel:            node.Data.Head.Level - 1,
+				NodeParentGeneration: node.Data.Head.Generation,
 			})
 			if cbs.PreKeyPointer != nil {
 				if err := cbs.PreKeyPointer(itemPath, item); err != nil {
@@ -405,9 +435,10 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key, uint32) int) (TreePath,
 		TreeID: treeRoot.TreeID,
 		Nodes: []TreePathElem{
 			{
-				ItemIdx:   -1,
-				NodeAddr:  treeRoot.RootNode,
-				NodeLevel: treeRoot.Level,
+				ItemIdx:              -1,
+				NodeAddr:             treeRoot.RootNode,
+				NodeLevel:            treeRoot.Level,
+				NodeParentGeneration: treeRoot.Generation,
 			},
 		},
 	}
@@ -415,7 +446,7 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key, uint32) int) (TreePath,
 		if path.Node(-1).NodeAddr == 0 {
 			return TreePath{}, nil, iofs.ErrNotExist
 		}
-		node, err := fs.readNodeAtLevel(path.Node(-1).NodeAddr, path.Node(-1).NodeLevel)
+		node, err := fs.ReadNode(path)
 		if err != nil {
 			return TreePath{}, nil, err
 		}
@@ -438,9 +469,10 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key, uint32) int) (TreePath,
 				return TreePath{}, nil, iofs.ErrNotExist
 			}
 			path = path.Append(TreePathElem{
-				ItemIdx:   lastGood,
-				NodeAddr:  node.Data.BodyInternal[lastGood].BlockPtr,
-				NodeLevel: node.Data.Head.Level - 1,
+				ItemIdx:              lastGood,
+				NodeAddr:             node.Data.BodyInternal[lastGood].BlockPtr,
+				NodeLevel:            node.Data.Head.Level - 1,
+				NodeParentGeneration: node.Data.Head.Generation,
 			})
 		} else {
 			// leaf node
@@ -484,7 +516,7 @@ func (fs *FS) prev(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	path.Node(-1).ItemIdx--
 	if path.Node(-1).NodeAddr != 0 {
 		if node.Addr != path.Node(-2).NodeAddr {
-			node, err = fs.readNodeAtLevel(path.Node(-2).NodeAddr, path.Node(-2).NodeLevel)
+			node, err = fs.ReadNode(path.Parent())
 			if err != nil {
 				return TreePath{}, nil, err
 			}
@@ -494,16 +526,17 @@ func (fs *FS) prev(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	// go down
 	for path.Node(-1).NodeAddr != 0 {
 		if node.Addr != path.Node(-1).NodeAddr {
-			node, err = fs.readNodeAtLevel(path.Node(-1).NodeAddr, path.Node(-1).NodeLevel)
+			node, err = fs.ReadNode(path)
 			if err != nil {
 				return TreePath{}, nil, err
 			}
 		}
 		if node.Data.Head.Level > 0 {
 			path = path.Append(TreePathElem{
-				ItemIdx:   len(node.Data.BodyInternal) - 1,
-				NodeAddr:  node.Data.BodyInternal[len(node.Data.BodyInternal)-1].BlockPtr,
-				NodeLevel: node.Data.Head.Level - 1,
+				ItemIdx:              len(node.Data.BodyInternal) - 1,
+				NodeAddr:             node.Data.BodyInternal[len(node.Data.BodyInternal)-1].BlockPtr,
+				NodeLevel:            node.Data.Head.Level - 1,
+				NodeParentGeneration: node.Data.Head.Generation,
 			})
 		} else {
 			path = path.Append(TreePathElem{
@@ -513,7 +546,7 @@ func (fs *FS) prev(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	}
 	// return
 	if node.Addr != path.Node(-2).NodeAddr {
-		node, err = fs.readNodeAtLevel(path.Node(-2).NodeAddr, path.Node(-2).NodeLevel)
+		node, err = fs.ReadNode(path.Parent())
 		if err != nil {
 			return TreePath{}, nil, err
 		}
@@ -527,7 +560,7 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 
 	// go up
 	if node.Addr != path.Node(-2).NodeAddr {
-		node, err = fs.readNodeAtLevel(path.Node(-2).NodeAddr, path.Node(-2).NodeLevel)
+		node, err = fs.ReadNode(path.Parent())
 		if err != nil {
 			return TreePath{}, nil, err
 		}
@@ -539,18 +572,18 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 			return TreePath{}, nil, nil
 		}
 		if node.Addr != path.Node(-2).NodeAddr {
-			node, err = fs.readNodeAtLevel(path.Node(-2).NodeAddr, path.Node(-2).NodeLevel)
+			node, err = fs.ReadNode(path.Parent())
 			if err != nil {
 				return TreePath{}, nil, err
 			}
 			path.Node(-2).NodeLevel = node.Data.Head.Level
 		}
 	}
-	// go left
+	// go right
 	path.Node(-1).ItemIdx++
 	if path.Node(-1).NodeAddr != 0 {
 		if node.Addr != path.Node(-2).NodeAddr {
-			node, err = fs.readNodeAtLevel(path.Node(-2).NodeAddr, path.Node(-2).NodeLevel)
+			node, err = fs.ReadNode(path.Parent())
 			if err != nil {
 				return TreePath{}, nil, err
 			}
@@ -560,7 +593,7 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	// go down
 	for path.Node(-1).NodeAddr != 0 {
 		if node.Addr != path.Node(-1).NodeAddr {
-			node, err = fs.readNodeAtLevel(path.Node(-1).NodeAddr, path.Node(-1).NodeLevel)
+			node, err = fs.ReadNode(path)
 			if err != nil {
 				return TreePath{}, nil, err
 			}
@@ -568,9 +601,10 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 		}
 		if node.Data.Head.Level > 0 {
 			path = path.Append(TreePathElem{
-				ItemIdx:   0,
-				NodeAddr:  node.Data.BodyInternal[len(node.Data.BodyInternal)-1].BlockPtr,
-				NodeLevel: node.Data.Head.Level - 1,
+				ItemIdx:              0,
+				NodeAddr:             node.Data.BodyInternal[len(node.Data.BodyInternal)-1].BlockPtr,
+				NodeLevel:            node.Data.Head.Level - 1,
+				NodeParentGeneration: node.Data.Head.Generation,
 			})
 		} else {
 			path = path.Append(TreePathElem{
@@ -580,7 +614,7 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	}
 	// return
 	if node.Addr != path.Node(-2).NodeAddr {
-		node, err = fs.readNodeAtLevel(path.Node(-2).NodeAddr, path.Node(-2).NodeLevel)
+		node, err = fs.ReadNode(path.Parent())
 		if err != nil {
 			return TreePath{}, nil, err
 		}
@@ -588,6 +622,7 @@ func (fs *FS) next(path TreePath, node *diskio.Ref[btrfsvol.LogicalAddr, Node])
 	return path, node, nil
 }
 
+// TreeSearch implements the 'Trees' interface.
 func (fs *FS) TreeSearch(treeID ObjID, fn func(Key, uint32) int) (Item, error) {
 	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
@@ -600,12 +635,14 @@ func (fs *FS) TreeSearch(treeID ObjID, fn func(Key, uint32) int) (Item, error) {
 	return node.Data.BodyLeaf[path.Node(-1).ItemIdx], nil
 }
 
+// KeySearch returns a comparator suitable to be passed to TreeSearch.
 func KeySearch(fn func(Key) int) func(Key, uint32) int {
 	return func(key Key, _ uint32) int {
 		return fn(key)
 	}
 }
 
+// TreeLookup implements the 'Trees' interface.
 func (fs *FS) TreeLookup(treeID ObjID, key Key) (Item, error) {
 	item, err := fs.TreeSearch(treeID, KeySearch(key.Cmp))
 	if err != nil {
@@ -614,6 +651,7 @@ func (fs *FS) TreeLookup(treeID ObjID, key Key) (Item, error) {
 	return item, err
 }
 
+// TreeSearchAll implements the 'Trees' interface.
 func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key, uint32) int) ([]Item, error) {
 	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
diff --git a/lib/btrfs/types_node.go b/lib/btrfs/types_node.go
index 3a83ee7..e7be341 100644
--- a/lib/btrfs/types_node.go
+++ b/lib/btrfs/types_node.go
@@ -13,6 +13,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfssum"
 	"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"
 	"git.lukeshu.com/btrfs-progs-ng/lib/fmtutil"
 )
@@ -378,7 +379,15 @@ func (node *Node) LeafFreeSpace() uint32 {
 
 var ErrNotANode = errors.New("does not look like a node")
 
-func ReadNode[Addr ~int64](fs diskio.File[Addr], sb Superblock, addr Addr, laddrCB func(btrfsvol.LogicalAddr) error) (*diskio.Ref[Addr, Node], error) {
+type NodeExpectations struct {
+	LAddr containers.Optional[btrfsvol.LogicalAddr]
+	// Things knowable from the parent.
+	Level         containers.Optional[uint8]
+	MaxGeneration containers.Optional[Generation]
+	Owner         containers.Optional[ObjID]
+}
+
+func ReadNode[Addr ~int64](fs diskio.File[Addr], sb Superblock, addr Addr, exp NodeExpectations) (*diskio.Ref[Addr, Node], error) {
 	nodeBuf := make([]byte, sb.NodeSize)
 	if _, err := fs.ReadAt(nodeBuf, addr); err != nil {
 		return nil, err
@@ -414,10 +423,21 @@ func ReadNode[Addr ~int64](fs diskio.File[Addr], sb Superblock, addr Addr, laddr
 			addr, stored, calced)
 	}
 
-	if laddrCB != nil {
-		if err := laddrCB(nodeRef.Data.Head.Addr); err != nil {
-			return nodeRef, fmt.Errorf("btrfs.ReadNode: node@%v: %w", addr, err)
-		}
+	if exp.LAddr.OK && nodeRef.Data.Head.Addr != exp.LAddr.Val {
+		return nodeRef, fmt.Errorf("btrfs.ReadNode: node@%v: read from laddr=%v but claims to be at laddr=%v",
+			addr, exp.LAddr.Val, nodeRef.Data.Head.Addr)
+	}
+	if exp.Level.OK && nodeRef.Data.Head.Level != exp.Level.Val {
+		return nodeRef, fmt.Errorf("btrfs.ReadNode: node@%v: expected level=%v but claims to be level=%v",
+			addr, exp.Level.Val, nodeRef.Data.Head.Level)
+	}
+	if exp.MaxGeneration.OK && nodeRef.Data.Head.Generation > exp.MaxGeneration.Val {
+		return nodeRef, fmt.Errorf("btrfs.ReadNode: node@%v: expected generation<=%v but claims to be generation=%v",
+			addr, exp.MaxGeneration.Val, nodeRef.Data.Head.Generation)
+	}
+	if exp.Owner.OK && nodeRef.Data.Head.Owner != exp.Owner.Val {
+		return nodeRef, fmt.Errorf("btrfs.ReadNode: node@%v: expected owner=%v but claims to have owner=%v",
+			addr, exp.Owner.Val, nodeRef.Data.Head.Owner)
 	}
 
 	// parse (main)
@@ -430,30 +450,3 @@ func ReadNode[Addr ~int64](fs diskio.File[Addr], sb Superblock, addr Addr, laddr
 
 	return nodeRef, nil
 }
-
-func (fs *FS) ReadNode(addr btrfsvol.LogicalAddr) (*diskio.Ref[btrfsvol.LogicalAddr, Node], error) {
-	sb, err := fs.Superblock()
-	if err != nil {
-		return nil, fmt.Errorf("btrfs.FS.ReadNode: %w", err)
-	}
-
-	return ReadNode[btrfsvol.LogicalAddr](fs, *sb, addr, func(claimAddr btrfsvol.LogicalAddr) error {
-		if claimAddr != addr {
-			return fmt.Errorf("read from laddr=%v but claims to be at laddr=%v",
-				addr, claimAddr)
-		}
-		return nil
-	})
-}
-
-func (fs *FS) readNodeAtLevel(addr btrfsvol.LogicalAddr, expLevel uint8) (*diskio.Ref[btrfsvol.LogicalAddr, Node], error) {
-	node, err := fs.ReadNode(addr)
-	if err != nil {
-		return node, err
-	}
-	if node.Data.Head.Level != expLevel {
-		return node, fmt.Errorf("btrfs.FS.ReadNode: node@%v: expected level %v but has level %v",
-			node.Addr, expLevel, node.Data.Head.Level)
-	}
-	return node, nil
-}