From 7c0e009092f593ca0990a50de8db3f81c112e58e Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sun, 18 Dec 2022 00:34:06 -0700
Subject: rebuildnodes: Reduce scope of the orphans set

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index ef50653..d197bf8 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -37,7 +37,6 @@ type Rebuilder struct {
 	graph        graph.Graph
 	uuidMap      uuidmap.UUIDMap
 	csums        containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
-	orphans      containers.Set[btrfsvol.LogicalAddr]
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
 	key2leaf     containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr]
 
@@ -65,7 +64,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Indexing orphans...")
-	orphans, leaf2orphans, key2leaf, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
+	leaf2orphans, key2leaf, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +78,6 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		graph:        *scanData.nodeGraph,
 		uuidMap:      *scanData.uuidMap,
 		csums:        *csums,
-		orphans:      orphans,
 		leaf2orphans: leaf2orphans,
 		key2leaf:     *key2leaf,
 
-- 
cgit v1.2.3-2-g168b


From abd7485b9fb78db233839effcd06e69d44da3859 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 20 Dec 2022 18:42:24 -0700
Subject: rebuildnodes: Don't export `Rebuilder`

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 24 +++++++++++-----------
 1 file changed, 12 insertions(+), 12 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index d197bf8..e441abd 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -26,7 +26,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
-type Rebuilder struct {
+type rebuilder struct {
 	raw   *btrfs.FS
 	inner interface {
 		btrfstree.TreeOperator
@@ -70,7 +70,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Rebuilding node tree...")
-	o := &Rebuilder{
+	o := &rebuilder{
 		raw:   fs,
 		inner: btrfsutil.NewBrokenTrees(ctx, fs),
 		sb:    *sb,
@@ -90,13 +90,13 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	return o.augments, nil
 }
 
-func (o *Rebuilder) ioErr(ctx context.Context, err error) {
+func (o *rebuilder) ioErr(ctx context.Context, err error) {
 	err = fmt.Errorf("should not happen: i/o error: %w", err)
 	dlog.Error(ctx, err)
 	panic(err)
 }
 
-func (o *Rebuilder) rebuild(ctx context.Context) error {
+func (o *rebuilder) rebuild(ctx context.Context) error {
 	passNum := 0
 	dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
 	o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
@@ -151,7 +151,7 @@ func (o *Rebuilder) rebuild(ctx context.Context) error {
 	return nil
 }
 
-func (o *Rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances []map[btrfsvol.LogicalAddr]int) containers.Set[btrfsvol.LogicalAddr] {
+func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances []map[btrfsvol.LogicalAddr]int) containers.Set[btrfsvol.LogicalAddr] {
 	distances := make(map[btrfsvol.LogicalAddr]int)
 	generations := make(map[btrfsvol.LogicalAddr]btrfsprim.Generation)
 	lists := make([]containers.Set[btrfsvol.LogicalAddr], len(listsWithDistances))
@@ -294,7 +294,7 @@ func treeDistance(fs _NodeFile, tree, leaf btrfsprim.ObjID) (int, bool) {
 	}
 }
 
-func (o *Rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
+func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
 	choicesWithDist := make(map[btrfsvol.LogicalAddr]int)
 	for choice := range choices {
 		if dist, ok := treeDistance(o.uuidMap, treeID, o.graph.Nodes[choice].Owner); ok {
@@ -312,12 +312,12 @@ func (o *Rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, cho
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 // fsErr implements rebuildCallbacks.
-func (o *Rebuilder) fsErr(ctx context.Context, e error) {
+func (o *rebuilder) fsErr(ctx context.Context, e error) {
 	dlog.Errorf(ctx, "filesystem error: %v", e)
 }
 
 // want implements rebuildCallbacks.
-func (o *Rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
+func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -345,7 +345,7 @@ func (o *Rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrf
 }
 
 // wantOff implements rebuildCallbacks.
-func (o *Rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
+func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -371,7 +371,7 @@ func (o *Rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 }
 
 // wantFunc implements rebuildCallbacks.
-func (o *Rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
+func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -415,7 +415,7 @@ func (o *Rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 // func implements rebuildCallbacks.
 //
 // interval is [beg, end)
-func (o *Rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
+func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
 	for beg < end {
 		// check if we already have it
 		if run, err := btrfs.LookupCSum(o.inner, o.sb.ChecksumType, beg); err == nil {
@@ -460,7 +460,7 @@ func (o *Rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 }
 
 // wantFileExt implements rebuildCallbacks.
-func (o *Rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
+func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
 	min := btrfsprim.Key{
 		ObjectID: ino,
 		ItemType: btrfsitem.EXTENT_DATA_KEY,
-- 
cgit v1.2.3-2-g168b


From cb93353c360258569e4b91e164136e321707470a Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 13:53:04 -0700
Subject: rebuildnodes: Fix a missing word in a comment

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index e441abd..e0ca7a0 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -590,8 +590,8 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 							continue
 						}
 						itemEnd := itemBeg + itemSize
-						// We're being and "wanting" any extent that has any overlap with the
-						// gap.  But maybe instead we sould only want extents that are
+						// We're being greedy and "wanting" any extent that has any overlap with
+						// the gap.  But maybe instead we sould only want extents that are
 						// *entirely* within the gap.  I'll have to run it on real filesystems
 						// to see what works better.
 						//
-- 
cgit v1.2.3-2-g168b


From 30be1eacbe4ff253c82c6e6163234dc20b4de550 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 05:30:25 -0700
Subject: rebuildnodes: Refactor existing key-io code in to a sub-package

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 118 +++++++++------------
 1 file changed, 53 insertions(+), 65 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index e0ca7a0..254cfee 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -20,6 +20,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildmappings"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
@@ -38,7 +39,7 @@ type rebuilder struct {
 	uuidMap      uuidmap.UUIDMap
 	csums        containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
-	key2leaf     containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr]
+	keyIO        keyio.Handle
 
 	augments map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
 
@@ -64,7 +65,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Indexing orphans...")
-	leaf2orphans, key2leaf, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
+	leaf2orphans, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +80,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		uuidMap:      *scanData.uuidMap,
 		csums:        *csums,
 		leaf2orphans: leaf2orphans,
-		key2leaf:     *key2leaf,
+		keyIO:        *scanData.keyIO,
 
 		augments: make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]),
 	}
@@ -335,10 +336,10 @@ func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrf
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.key2leaf.Subrange(
-		func(k keyAndTree, _ btrfsvol.LogicalAddr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(_ keyAndTree, v btrfsvol.LogicalAddr) bool {
-			wants.InsertFrom(o.leaf2orphans[v])
+	o.keyIO.Keys.Subrange(
+		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
+		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			wants.InsertFrom(o.leaf2orphans[v.Node])
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -361,10 +362,10 @@ func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.key2leaf.Subrange(
-		func(k keyAndTree, _ btrfsvol.LogicalAddr) int { return tgt.Cmp(k.Key) },
-		func(_ keyAndTree, v btrfsvol.LogicalAddr) bool {
-			wants.InsertFrom(o.leaf2orphans[v])
+	o.keyIO.Keys.Subrange(
+		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { return tgt.Cmp(k.Key) },
+		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			wants.InsertFrom(o.leaf2orphans[v.Node])
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -392,20 +393,15 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v +func", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.key2leaf.Subrange(
-		func(k keyAndTree, _ btrfsvol.LogicalAddr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(k keyAndTree, v btrfsvol.LogicalAddr) bool {
-			nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](o.raw, o.sb, v, btrfstree.NodeExpectations{
-				LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: v},
-				Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: o.graph.Nodes[v].Generation},
-			})
-			if err != nil {
-				o.ioErr(ctx, err)
+	o.keyIO.Keys.Subrange(
+		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
+		func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			itemBody, ok := o.keyIO.ReadItem(v)
+			if !ok {
+				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
 			}
-			for _, item := range nodeRef.Data.BodyLeaf {
-				if k.Key == item.Key && fn(item.Body) {
-					wants.InsertFrom(o.leaf2orphans[v])
-				}
+			if fn(itemBody) {
+				wants.InsertFrom(o.leaf2orphans[v.Node])
 			}
 			return true
 		})
@@ -441,18 +437,18 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 				continue
 			}
 			run := rbNode.Value.Sums
-			key := keyAndTree{
+			key := keyio.KeyAndTree{
 				Key:    rbNode.Value.Key,
 				TreeID: btrfsprim.CSUM_TREE_OBJECTID,
 			}
-			leaf, ok := o.key2leaf.Load(key)
+			itemPtr, ok := o.keyIO.Keys.Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
 				// to be in some Node, and if we didn't find it from
 				// btrfs.LookupCSum(), then that Node must be an orphan.
 				panic(fmt.Errorf("should not happen: no orphan contains %v", key.Key))
 			}
-			o.wantAugment(ctx, key.TreeID, o.leaf2orphans[leaf])
+			o.wantAugment(ctx, key.TreeID, o.leaf2orphans[itemPtr.Node])
 
 			beg = run.Addr.Add(run.Size())
 		}
@@ -558,8 +554,8 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		}
 		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("file extent for tree=%v inode=%v bytes [%v, %v)", treeID, ino, gap.Beg, gap.End))
 		wants := make(containers.Set[btrfsvol.LogicalAddr])
-		o.key2leaf.Subrange(
-			func(k keyAndTree, _ btrfsvol.LogicalAddr) int {
+		o.keyIO.Keys.Subrange(
+			func(k keyio.KeyAndTree, _ keyio.ItemPtr) int {
 				switch {
 				case min.Cmp(k.Key) < 0:
 					return 1
@@ -569,45 +565,37 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 					return 0
 				}
 			},
-			func(k keyAndTree, v btrfsvol.LogicalAddr) bool {
-				nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](o.raw, o.sb, v, btrfstree.NodeExpectations{
-					LAddr:      containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: v},
-					Generation: containers.Optional[btrfsprim.Generation]{OK: true, Val: o.graph.Nodes[v].Generation},
-				})
-				if err != nil {
-					o.ioErr(ctx, fmt.Errorf("error reading previously read node@%v: %w", v, err))
+			func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+				itemBody, ok := o.keyIO.ReadItem(v)
+				if !ok {
+					o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
 				}
-				for _, item := range nodeRef.Data.BodyLeaf {
-					if k.Key != item.Key {
-						continue
+				switch itemBody := itemBody.(type) {
+				case btrfsitem.FileExtent:
+					itemBeg := int64(k.Offset)
+					itemSize, err := itemBody.Size()
+					if err != nil {
+						o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, k, err))
+						break
 					}
-					switch itemBody := item.Body.(type) {
-					case btrfsitem.FileExtent:
-						itemBeg := int64(item.Key.Offset)
-						itemSize, err := itemBody.Size()
-						if err != nil {
-							o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, item.Key, err))
-							continue
-						}
-						itemEnd := itemBeg + itemSize
-						// We're being greedy and "wanting" any extent that has any overlap with
-						// the gap.  But maybe instead we sould only want extents that are
-						// *entirely* within the gap.  I'll have to run it on real filesystems
-						// to see what works better.
-						//
-						// TODO(lukeshu): Re-evaluate whether being greedy here is the right
-						// thing.
-						if itemEnd > gap.Beg && itemBeg < gap.End {
-							wants.InsertFrom(o.leaf2orphans[v])
-						}
-					case btrfsitem.Error:
-						o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, item.Key, itemBody.Err))
-					default:
-						// This is a panic because the item decoder should not emit EXTENT_DATA
-						// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-						// this code also being updated.
-						panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", itemBody))
+					itemEnd := itemBeg + itemSize
+					// We're being greedy and "wanting" any extent that has any overlap with
+					// the gap.  But maybe instead we sould only want extents that are
+					// *entirely* within the gap.  I'll have to run it on real filesystems
+					// to see what works better.
+					//
+					// TODO(lukeshu): Re-evaluate whether being greedy here is the right
+					// thing.
+					if itemEnd > gap.Beg && itemBeg < gap.End {
+						wants.InsertFrom(o.leaf2orphans[v.Node])
 					}
+				case btrfsitem.Error:
+					o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, k, itemBody.Err))
+				default:
+					// This is a panic because the item decoder should not emit EXTENT_DATA
+					// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
+					// this code also being updated.
+					panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", itemBody))
 				}
 				return true
 			})
-- 
cgit v1.2.3-2-g168b


From 7620e1948a4d1e17458d8cfc9ed306bb774a3274 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 21 Dec 2022 04:30:56 -0700
Subject: rebuildnodes: Migrate to the new rebuilt-btrees system

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 396 +++++++++++++--------
 1 file changed, 245 insertions(+), 151 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 254cfee..2ddce88 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"sort"
+	"time"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -19,30 +20,24 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildmappings"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 type rebuilder struct {
-	raw   *btrfs.FS
-	inner interface {
-		btrfstree.TreeOperator
-		Augment(treeID btrfsprim.ObjID, nodeAddr btrfsvol.LogicalAddr) ([]btrfsprim.Key, error)
-	}
-	sb btrfstree.Superblock
-
-	graph        graph.Graph
-	uuidMap      uuidmap.UUIDMap
-	csums        containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
-	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
-	keyIO        keyio.Handle
+	sb      btrfstree.Superblock
+	rebuilt *btrees.RebuiltTrees
 
-	augments map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
+	graph graph.Graph
+	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
+	keyIO keyio.Handle
 
+	curKey          keyio.KeyAndTree
+	queue           []keyio.KeyAndTree
 	pendingAugments map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int
 }
 
@@ -64,31 +59,21 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		csums = new(containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum])
 	}
 
-	dlog.Info(ctx, "Indexing orphans...")
-	leaf2orphans, err := indexOrphans(ctx, fs, *sb, *scanData.nodeGraph)
-	if err != nil {
-		return nil, err
-	}
-
 	dlog.Info(ctx, "Rebuilding node tree...")
 	o := &rebuilder{
-		raw:   fs,
-		inner: btrfsutil.NewBrokenTrees(ctx, fs),
-		sb:    *sb,
-
-		graph:        *scanData.nodeGraph,
-		uuidMap:      *scanData.uuidMap,
-		csums:        *csums,
-		leaf2orphans: leaf2orphans,
-		keyIO:        *scanData.keyIO,
+		sb: *sb,
 
-		augments: make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]),
+		graph: *scanData.nodeGraph,
+		csums: *csums,
+		keyIO: *scanData.keyIO,
 	}
+	o.rebuilt = btrees.NewRebuiltTrees(*sb, *scanData.nodeGraph, *scanData.keyIO,
+		o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID)
 	if err := o.rebuild(ctx); err != nil {
 		return nil, err
 	}
 
-	return o.augments, nil
+	return o.rebuilt.ListRoots(), nil
 }
 
 func (o *rebuilder) ioErr(ctx context.Context, err error) {
@@ -97,61 +82,153 @@ func (o *rebuilder) ioErr(ctx context.Context, err error) {
 	panic(err)
 }
 
+type rebuildStats struct {
+	PassNum int
+	Task    string
+	N, D    int
+}
+
+func (s rebuildStats) String() string {
+	pct := 100
+	if s.D > 0 {
+		pct = int(100 * float64(s.N) / float64(s.D))
+	}
+	return fmt.Sprintf("... pass %d: %s %v%% (%v/%v)",
+		s.PassNum, s.Task, pct, s.N, s.D)
+}
+
 func (o *rebuilder) rebuild(ctx context.Context) error {
 	passNum := 0
 	dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
+	// Seed the queue
 	o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
-	btrfsutil.WalkAllTrees(ctx, o.inner, btrfsutil.WalkAllTreesHandler{
-		Err: func(*btrfsutil.WalkError) {},
-		TreeWalkHandler: btrfstree.TreeWalkHandler{
-			Item: func(path btrfstree.TreePath, item btrfstree.Item) error {
-				handleItem(o, ctx, path[0].FromTree, item)
-				return nil
-			},
-		},
-	})
-	for len(o.pendingAugments) > 0 {
-		// Apply the augments, keeping track of what keys are added to what tree.
-		dlog.Infof(ctx, "... pass %d: augmenting trees to add implied items", passNum)
-		newKeys := make(map[btrfsprim.ObjID][]btrfsprim.Key)
-		for _, treeID := range maps.SortedKeys(o.pendingAugments) {
-			dlog.Infof(ctx, "... ... augmenting tree %v:", treeID)
-			treeAugments := o.resolveTreeAugments(ctx, o.pendingAugments[treeID])
-			for _, nodeAddr := range maps.SortedKeys(treeAugments) {
-				added, err := o.inner.Augment(treeID, nodeAddr)
-				if err != nil {
-					dlog.Errorf(ctx, "error augmenting: %v", err)
-					continue
-				}
-				newKeys[treeID] = append(newKeys[treeID], added...)
-
-				set := o.augments[treeID]
-				if set == nil {
-					set = make(containers.Set[btrfsvol.LogicalAddr])
-					o.augments[treeID] = set
-				}
-				set.Insert(nodeAddr)
+	o.rebuilt.AddTree(ctx, btrfsprim.ROOT_TREE_OBJECTID)
+	o.rebuilt.AddTree(ctx, btrfsprim.CHUNK_TREE_OBJECTID)
+	o.rebuilt.AddTree(ctx, btrfsprim.TREE_LOG_OBJECTID)
+	o.rebuilt.AddTree(ctx, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
+	for {
+		// Handle items in the queue
+		queue := o.queue
+		o.queue = nil
+		progressWriter := textui.NewProgress[rebuildStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+		queueProgress := func(done int) {
+			progressWriter.Set(rebuildStats{
+				PassNum: passNum,
+				Task:    "processing item queue",
+				N:       done,
+				D:       len(queue),
+			})
+		}
+		for i, key := range queue {
+			queueProgress(i)
+			o.curKey = key
+			itemBody, ok := o.rebuilt.Load(ctx, key.TreeID, key.Key)
+			if !ok {
+				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
+			}
+			handleItem(o, ctx, key.TreeID, btrfstree.Item{
+				Key:  key.Key,
+				Body: itemBody,
+			})
+			if key.ItemType == btrfsitem.ROOT_ITEM_KEY {
+				o.rebuilt.AddTree(ctx, key.ObjectID)
 			}
 		}
-		// Clear the list of pending augments.
+		queueProgress(len(queue))
+		progressWriter.Done()
+
+		// Check if we can bail
+		if len(o.queue) == 0 && len(o.pendingAugments) == 0 {
+			break
+		}
+
+		// Apply augments that were requested while handling items from the queue
+		resolvedAugments := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], len(o.pendingAugments))
+		numAugments := 0
+		for _, treeID := range maps.SortedKeys(o.pendingAugments) {
+			resolvedAugments[treeID] = o.resolveTreeAugments(ctx, o.pendingAugments[treeID])
+			numAugments += len(resolvedAugments[treeID])
+		}
 		o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
-		passNum++
-		// Call handleItem() for each of the added keys.
-		dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
-		for _, treeID := range maps.SortedKeys(newKeys) {
-			for _, key := range newKeys[treeID] {
-				item, err := o.inner.TreeLookup(treeID, key)
-				if err != nil {
-					o.ioErr(ctx, fmt.Errorf("error looking up already-inserted item: tree=%v key=%v: %w",
-						treeID, key, err))
-				}
-				handleItem(o, ctx, treeID, item)
+		progressWriter = textui.NewProgress[rebuildStats](ctx, dlog.LogLevelInfo, 1*time.Second)
+		numAugmented := 0
+		augmentProgress := func() {
+			progressWriter.Set(rebuildStats{
+				PassNum: passNum,
+				Task:    "applying augments",
+				N:       numAugmented,
+				D:       numAugments,
+			})
+		}
+		for _, treeID := range maps.SortedKeys(resolvedAugments) {
+			for _, nodeAddr := range maps.SortedKeys(resolvedAugments[treeID]) {
+				augmentProgress()
+				o.rebuilt.AddRoot(ctx, treeID, nodeAddr)
+				numAugmented++
 			}
 		}
+		augmentProgress()
+		progressWriter.Done()
+
+		passNum++
+		dlog.Infof(ctx, "... pass %d: scanning for implied items", passNum)
 	}
 	return nil
 }
 
+func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key) {
+	o.queue = append(o.queue, keyio.KeyAndTree{
+		TreeID: tree,
+		Key:    key,
+	})
+}
+
+func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool) {
+	key, ok := o._want(ctx, btrfsprim.ROOT_TREE_OBJECTID, tree, btrfsitem.ROOT_ITEM_KEY)
+	if !ok {
+		o.queue = append(o.queue, o.curKey)
+		return btrfsitem.Root{}, false
+	}
+	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.ROOT_TREE_OBJECTID, key)
+	if !ok {
+		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
+	}
+	switch itemBody := itemBody.(type) {
+	case btrfsitem.Root:
+		return itemBody, true
+	case btrfsitem.Error:
+		o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", btrfsprim.ROOT_TREE_OBJECTID, key, itemBody.Err))
+		return btrfsitem.Root{}, false
+	default:
+		// This is a panic because the item decoder should not emit ROOT_ITEM items as anything but
+		// btrfsitem.Root or btrfsitem.Error without this code also being updated.
+		panic(fmt.Errorf("should not happen: ROOT_ITEM item has unexpected type: %T", itemBody))
+	}
+}
+
+func (o *rebuilder) cbLookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool) {
+	key := btrfsitem.UUIDToKey(uuid)
+	if ok := o._wantOff(ctx, btrfsprim.UUID_TREE_OBJECTID, key.ObjectID, key.ItemType, key.Offset); !ok {
+		o.queue = append(o.queue, o.curKey)
+		return 0, false
+	}
+	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.UUID_TREE_OBJECTID, key)
+	if !ok {
+		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
+	}
+	switch itemBody := itemBody.(type) {
+	case btrfsitem.UUIDMap:
+		return itemBody.ObjID, true
+	case btrfsitem.Error:
+		o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", btrfsprim.UUID_TREE_OBJECTID, key, itemBody.Err))
+		return 0, false
+	default:
+		// This is a panic because the item decoder should not emit UUID_SUBVOL items as anything but
+		// btrfsitem.UUIDMap or btrfsitem.Error without this code also being updated.
+		panic(fmt.Errorf("should not happen: UUID_SUBVOL item has unexpected type: %T", itemBody))
+	}
+}
+
 func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances []map[btrfsvol.LogicalAddr]int) containers.Set[btrfsvol.LogicalAddr] {
 	distances := make(map[btrfsvol.LogicalAddr]int)
 	generations := make(map[btrfsvol.LogicalAddr]btrfsprim.Generation)
@@ -268,37 +345,10 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-// _NodeFile is a subset of btrfstree.NodeFile.
-type _NodeFile interface {
-	ParentTree(btrfsprim.ObjID) (btrfsprim.ObjID, bool)
-}
-
-func treeDistance(fs _NodeFile, tree, leaf btrfsprim.ObjID) (int, bool) {
-	dist := 0
-	for {
-		if tree == leaf {
-			return dist, true
-		}
-
-		parentTree, ok := fs.ParentTree(tree)
-		if !ok {
-			// Failed to look up parent info.
-			return 0, false
-		}
-		if parentTree == 0 {
-			// End of the line.
-			return 0, false
-		}
-
-		tree = parentTree
-		dist++
-	}
-}
-
 func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
 	choicesWithDist := make(map[btrfsvol.LogicalAddr]int)
 	for choice := range choices {
-		if dist, ok := treeDistance(o.uuidMap, treeID, o.graph.Nodes[choice].Owner); ok {
+		if dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner); ok {
 			choicesWithDist[choice] = dist
 		}
 	}
@@ -319,17 +369,25 @@ func (o *rebuilder) fsErr(ctx context.Context, e error) {
 
 // want implements rebuildCallbacks.
 func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
+	o._want(ctx, treeID, objID, typ)
+}
+func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) (key btrfsprim.Key, ok bool) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return btrfsprim.Key{}, false
+	}
+
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	if _, err := o.inner.TreeSearch(treeID, func(key btrfsprim.Key, _ uint32) int {
+	if key, ok := o.rebuilt.Search(ctx, treeID, func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
-	}); err == nil {
-		return
+	}); ok {
+		return key, true
 	}
 
 	// OK, we need to insert it
@@ -339,14 +397,23 @@ func (o *rebuilder) want(ctx context.Context, treeID btrfsprim.ObjID, objID btrf
 	o.keyIO.Keys.Subrange(
 		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
 		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.leaf2orphans[v.Node])
+			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
+	return btrfsprim.Key{}, false
 }
 
 // wantOff implements rebuildCallbacks.
 func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
+	o._wantOff(ctx, treeID, objID, typ, off)
+}
+func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) (ok bool) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return false
+	}
+
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
@@ -354,8 +421,8 @@ func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 		ItemType: typ,
 		Offset:   off,
 	}
-	if _, err := o.inner.TreeLookup(treeID, tgt); err == nil {
-		return
+	if _, ok := o.rebuilt.Load(ctx, treeID, tgt); ok {
+		return true
 	}
 
 	// OK, we need to insert it
@@ -365,26 +432,36 @@ func (o *rebuilder) wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID b
 	o.keyIO.Keys.Subrange(
 		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { return tgt.Cmp(k.Key) },
 		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.leaf2orphans[v.Node])
+			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
+	return false
 }
 
 // wantFunc implements rebuildCallbacks.
 func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(btrfsitem.Item) bool) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return
+	}
+
 	// check if we already have it
 
 	tgt := btrfsprim.Key{
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	items, _ := o.inner.TreeSearchAll(treeID, func(key btrfsprim.Key, _ uint32) int {
+	keys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	})
-	for _, item := range items {
-		if fn(item.Body) {
+	for _, itemKey := range keys {
+		itemBody, ok := o.rebuilt.Load(ctx, treeID, itemKey)
+		if !ok {
+			o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", itemKey))
+		}
+		if fn(itemBody) {
 			return
 		}
 	}
@@ -398,10 +475,10 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 		func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
 			itemBody, ok := o.keyIO.ReadItem(v)
 			if !ok {
-				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
+				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v at %v", k, v))
 			}
 			if fn(itemBody) {
-				wants.InsertFrom(o.leaf2orphans[v.Node])
+				wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			}
 			return true
 		})
@@ -412,35 +489,42 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 //
 // interval is [beg, end)
 func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
-	for beg < end {
-		// check if we already have it
-		if run, err := btrfs.LookupCSum(o.inner, o.sb.ChecksumType, beg); err == nil {
-			// we already have it
-			beg = run.Addr.Add(run.Size())
-		} else {
-			// we need to insert it
-			ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("csum for laddr=%v", beg))
-			rbNode := o.csums.Search(func(item btrfsinspect.SysExtentCSum) int {
-				switch {
-				case item.Sums.Addr > beg:
-					return -1
-				case item.Sums.Addr.Add(item.Sums.Size()) <= beg:
-					return 1
-				default:
-					return 0
-				}
+	if !o.rebuilt.AddTree(ctx, btrfsprim.CSUM_TREE_OBJECTID) {
+		o.queue = append(o.queue, o.curKey)
+		return
+	}
 
-			})
-			if rbNode == nil {
-				o.wantAugment(ctx, btrfsprim.CSUM_TREE_OBJECTID, nil) // log an error
-				beg += btrfssum.BlockSize
-				continue
-			}
-			run := rbNode.Value.Sums
-			key := keyio.KeyAndTree{
-				Key:    rbNode.Value.Key,
-				TreeID: btrfsprim.CSUM_TREE_OBJECTID,
+	for beg < end {
+		// Figure out which key we want.
+
+		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("csum for laddr=%v", beg))
+		rbNode := o.csums.Search(func(item btrfsinspect.SysExtentCSum) int {
+			switch {
+			case item.Sums.Addr > beg:
+				return -1
+			case item.Sums.Addr.Add(item.Sums.Size()) <= beg:
+				return 1
+			default:
+				return 0
 			}
+
+		})
+		if rbNode == nil {
+			o.wantAugment(ctx, btrfsprim.CSUM_TREE_OBJECTID, nil) // log an error
+			beg += btrfssum.BlockSize
+			continue
+		}
+		run := rbNode.Value.Sums
+		key := keyio.KeyAndTree{
+			Key:    rbNode.Value.Key,
+			TreeID: btrfsprim.CSUM_TREE_OBJECTID,
+		}
+
+		// Check if we already have it.
+
+		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
+		if _, ok := o.rebuilt.Search(ctx, key.TreeID, key.Key.Cmp); !ok {
+			// We need to insert it.
 			itemPtr, ok := o.keyIO.Keys.Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
@@ -448,15 +532,20 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 				// btrfs.LookupCSum(), then that Node must be an orphan.
 				panic(fmt.Errorf("should not happen: no orphan contains %v", key.Key))
 			}
-			o.wantAugment(ctx, key.TreeID, o.leaf2orphans[itemPtr.Node])
-
-			beg = run.Addr.Add(run.Size())
+			o.wantAugment(ctx, key.TreeID, o.rebuilt.LeafToRoots(ctx, key.TreeID, itemPtr.Node))
 		}
+
+		beg = run.Addr.Add(run.Size())
 	}
 }
 
 // wantFileExt implements rebuildCallbacks.
 func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
+	if !o.rebuilt.AddTree(ctx, treeID) {
+		o.queue = append(o.queue, o.curKey)
+		return
+	}
+
 	min := btrfsprim.Key{
 		ObjectID: ino,
 		ItemType: btrfsitem.EXTENT_DATA_KEY,
@@ -467,7 +556,7 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		ItemType: btrfsitem.EXTENT_DATA_KEY,
 		Offset:   uint64(size - 1),
 	}
-	exts, _ := o.inner.TreeSearchAll(treeID, func(key btrfsprim.Key, _ uint32) int {
+	extKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
 		switch {
 		case min.Cmp(key) < 0:
 			return 1
@@ -491,13 +580,18 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		Beg: 0,
 		End: size,
 	})
-	for _, ext := range exts {
-		switch extBody := ext.Body.(type) {
+	for _, extKey := range extKeys {
+		extBody, ok := o.rebuilt.Load(ctx, treeID, extKey)
+		if !ok {
+			o.ioErr(ctx, fmt.Errorf("could not look up already-inserted item: tree=%v key=%v",
+				treeID, extKey))
+		}
+		switch extBody := extBody.(type) {
 		case btrfsitem.FileExtent:
-			extBeg := int64(ext.Key.Offset)
+			extBeg := int64(extKey.Offset)
 			extSize, err := extBody.Size()
 			if err != nil {
-				o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, ext.Key, err))
+				o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, extKey, err))
 				continue
 			}
 			extEnd := extBeg + extSize
@@ -532,7 +626,7 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 				})
 			}
 		case btrfsitem.Error:
-			o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, ext.Key, extBody.Err))
+			o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, extKey, extBody.Err))
 		default:
 			// This is a panic because the item decoder should not emit EXTENT_DATA
 			// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
@@ -587,7 +681,7 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 					// TODO(lukeshu): Re-evaluate whether being greedy here is the right
 					// thing.
 					if itemEnd > gap.Beg && itemBeg < gap.End {
-						wants.InsertFrom(o.leaf2orphans[v.Node])
+						wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 					}
 				case btrfsitem.Error:
 					o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, k, itemBody.Err))
-- 
cgit v1.2.3-2-g168b


From a476f616f1e4367a212ef9b03434f593b39692e5 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 22 Dec 2022 02:56:56 -0700
Subject: rebuildnodes: Improve logging

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 2ddce88..e05399e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -146,6 +146,7 @@ func (o *rebuilder) rebuild(ctx context.Context) error {
 		resolvedAugments := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], len(o.pendingAugments))
 		numAugments := 0
 		for _, treeID := range maps.SortedKeys(o.pendingAugments) {
+			dlog.Infof(ctx, "... ... augments for tree %v:", treeID)
 			resolvedAugments[treeID] = o.resolveTreeAugments(ctx, o.pendingAugments[treeID])
 			numAugments += len(resolvedAugments[treeID])
 		}
@@ -337,7 +338,12 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances
 	}
 
 	for i, list := range lists {
-		dlog.Infof(ctx, "... ... ... %d: %v: %v", i, list.Intersection(ret).TakeOne(), maps.SortedKeys(list))
+		chose := list.Intersection(ret)
+		if len(chose) == 0 {
+			dlog.Infof(ctx, "... ... ... lists[%d]: chose (none) from %v", i, maps.SortedKeys(list))
+		} else {
+			dlog.Infof(ctx, "... ... ... lists[%d]: chose %v from %v", i, chose.TakeOne(), maps.SortedKeys(list))
+		}
 	}
 
 	return ret
-- 
cgit v1.2.3-2-g168b


From 3e22a30a7febd9890163966d01dd2b34656c9e39 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 22 Dec 2022 18:24:03 -0700
Subject: rebuildnodes: Handle the same node being added twice

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index e05399e..4e60632 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -352,16 +352,18 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, listsWithDistances
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, choices containers.Set[btrfsvol.LogicalAddr]) {
-	choicesWithDist := make(map[btrfsvol.LogicalAddr]int)
-	for choice := range choices {
-		if dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner); ok {
-			choicesWithDist[choice] = dist
-		}
-	}
-	if len(choicesWithDist) == 0 {
+	if len(choices) == 0 {
 		dlog.Errorf(ctx, "augment(tree=%v): could not find wanted item", treeID)
 		return
 	}
+	choicesWithDist := make(map[btrfsvol.LogicalAddr]int, len(choices))
+	for choice := range choices {
+		dist, ok := o.rebuilt.COWDistance(ctx, treeID, o.graph.Nodes[choice].Owner)
+		if !ok {
+			panic(fmt.Errorf("should not happen: .wantAugment called for tree=%v with invalid choice=%v", treeID, choice))
+		}
+		choicesWithDist[choice] = dist
+	}
 	dlog.Infof(ctx, "augment(tree=%v): %v", treeID, maps.SortedKeys(choicesWithDist))
 	o.pendingAugments[treeID] = append(o.pendingAugments[treeID], choicesWithDist)
 }
-- 
cgit v1.2.3-2-g168b


From fc6f8871c39d3a74153e4ce8fd440c256f65b650 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 16:03:27 -0700
Subject: rebuildnodes/btrees: Track the tree's root item offset

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 4e60632..180edab 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -184,11 +184,11 @@ func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key b
 	})
 }
 
-func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (item btrfsitem.Root, ok bool) {
+func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool) {
 	key, ok := o._want(ctx, btrfsprim.ROOT_TREE_OBJECTID, tree, btrfsitem.ROOT_ITEM_KEY)
 	if !ok {
 		o.queue = append(o.queue, o.curKey)
-		return btrfsitem.Root{}, false
+		return 0, btrfsitem.Root{}, false
 	}
 	itemBody, ok := o.rebuilt.Load(ctx, btrfsprim.ROOT_TREE_OBJECTID, key)
 	if !ok {
@@ -196,10 +196,10 @@ func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (ite
 	}
 	switch itemBody := itemBody.(type) {
 	case btrfsitem.Root:
-		return itemBody, true
+		return btrfsprim.Generation(key.Offset), itemBody, true
 	case btrfsitem.Error:
 		o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", btrfsprim.ROOT_TREE_OBJECTID, key, itemBody.Err))
-		return btrfsitem.Root{}, false
+		return 0, btrfsitem.Root{}, false
 	default:
 		// This is a panic because the item decoder should not emit ROOT_ITEM items as anything but
 		// btrfsitem.Root or btrfsitem.Error without this code also being updated.
-- 
cgit v1.2.3-2-g168b


From 51d6ccf1e5af4d88bfe6842b06fa1bc2d5cc732c Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 17:10:47 -0700
Subject: rebuildnodes: Have keyio.Handle always be a pointer, fuss with
 ScanDevices return type

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 180edab..fecf5e6 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -34,7 +34,7 @@ type rebuilder struct {
 
 	graph graph.Graph
 	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
-	keyIO keyio.Handle
+	keyIO *keyio.Handle
 
 	curKey          keyio.KeyAndTree
 	queue           []keyio.KeyAndTree
@@ -42,7 +42,7 @@ type rebuilder struct {
 }
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	scanData, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
@@ -63,11 +63,11 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	o := &rebuilder{
 		sb: *sb,
 
-		graph: *scanData.nodeGraph,
+		graph: nodeGraph,
 		csums: *csums,
-		keyIO: *scanData.keyIO,
+		keyIO: keyIO,
 	}
-	o.rebuilt = btrees.NewRebuiltTrees(*sb, *scanData.nodeGraph, *scanData.keyIO,
+	o.rebuilt = btrees.NewRebuiltTrees(*sb, nodeGraph, keyIO,
 		o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID)
 	if err := o.rebuild(ctx); err != nil {
 		return nil, err
-- 
cgit v1.2.3-2-g168b


From 6330b26a9f9c7769c48e2bc90e065457f8e138ab Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 16:48:42 -0700
Subject: rebuildnodes: Have the key index belong to the btree, and be smarter

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 45 +++++++++++-----------
 1 file changed, 22 insertions(+), 23 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index fecf5e6..982c27d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -42,7 +42,7 @@ type rebuilder struct {
 }
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+	nodeGraph, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
@@ -60,6 +60,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Rebuilding node tree...")
+	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 	o := &rebuilder{
 		sb: *sb,
 
@@ -402,9 +403,9 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.keyIO.Keys.Subrange(
-		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+	o.rebuilt.Keys(treeID).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
+		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
 			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
@@ -437,9 +438,9 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, objID
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.keyIO.Keys.Subrange(
-		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { return tgt.Cmp(k.Key) },
-		func(_ keyio.KeyAndTree, v keyio.ItemPtr) bool {
+	o.rebuilt.Keys(treeID).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int { return tgt.Cmp(k) },
+		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
 			wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
 			return true
 		})
@@ -478,9 +479,9 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 
 	ctx = dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key=%v +func", treeID, tgt))
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.keyIO.Keys.Subrange(
-		func(k keyio.KeyAndTree, _ keyio.ItemPtr) int { k.Key.Offset = 0; return tgt.Cmp(k.Key) },
-		func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+	o.rebuilt.Keys(treeID).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
+		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 			itemBody, ok := o.keyIO.ReadItem(v)
 			if !ok {
 				o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v at %v", k, v))
@@ -523,24 +524,22 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 			continue
 		}
 		run := rbNode.Value.Sums
-		key := keyio.KeyAndTree{
-			Key:    rbNode.Value.Key,
-			TreeID: btrfsprim.CSUM_TREE_OBJECTID,
-		}
+		key := rbNode.Value.Key
+		treeID := btrfsprim.CSUM_TREE_OBJECTID
 
 		// Check if we already have it.
 
 		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
-		if _, ok := o.rebuilt.Search(ctx, key.TreeID, key.Key.Cmp); !ok {
+		if _, ok := o.rebuilt.Search(ctx, treeID, key.Cmp); !ok {
 			// We need to insert it.
-			itemPtr, ok := o.keyIO.Keys.Load(key)
+			itemPtr, ok := o.rebuilt.Keys(btrfsprim.CSUM_TREE_OBJECTID).Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
 				// to be in some Node, and if we didn't find it from
 				// btrfs.LookupCSum(), then that Node must be an orphan.
-				panic(fmt.Errorf("should not happen: no orphan contains %v", key.Key))
+				panic(fmt.Errorf("should not happen: no orphan contains %v", key))
 			}
-			o.wantAugment(ctx, key.TreeID, o.rebuilt.LeafToRoots(ctx, key.TreeID, itemPtr.Node))
+			o.wantAugment(ctx, treeID, o.rebuilt.LeafToRoots(ctx, treeID, itemPtr.Node))
 		}
 
 		beg = run.Addr.Add(run.Size())
@@ -656,18 +655,18 @@ func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino
 		}
 		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("file extent for tree=%v inode=%v bytes [%v, %v)", treeID, ino, gap.Beg, gap.End))
 		wants := make(containers.Set[btrfsvol.LogicalAddr])
-		o.keyIO.Keys.Subrange(
-			func(k keyio.KeyAndTree, _ keyio.ItemPtr) int {
+		o.rebuilt.Keys(treeID).Subrange(
+			func(k btrfsprim.Key, _ keyio.ItemPtr) int {
 				switch {
-				case min.Cmp(k.Key) < 0:
+				case min.Cmp(k) < 0:
 					return 1
-				case max.Cmp(k.Key) > 0:
+				case max.Cmp(k) > 0:
 					return -1
 				default:
 					return 0
 				}
 			},
-			func(k keyio.KeyAndTree, v keyio.ItemPtr) bool {
+			func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 				itemBody, ok := o.keyIO.ReadItem(v)
 				if !ok {
 					o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
-- 
cgit v1.2.3-2-g168b


From 7ca8e6805164a5855c4af6ec5103c8e0cbcab89c Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 18:00:47 -0700
Subject: =?UTF-8?q?rebuildnodes:=20Move=20keyio.KeyAndTree=E2=86=92keyAndt?=
 =?UTF-8?q?ree?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 982c27d..00dbad5 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -28,6 +28,18 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
+type keyAndTree struct {
+	btrfsprim.Key
+	TreeID btrfsprim.ObjID
+}
+
+func (a keyAndTree) Cmp(b keyAndTree) int {
+	if d := a.Key.Cmp(b.Key); d != 0 {
+		return d
+	}
+	return containers.NativeCmp(a.TreeID, b.TreeID)
+}
+
 type rebuilder struct {
 	sb      btrfstree.Superblock
 	rebuilt *btrees.RebuiltTrees
@@ -36,8 +48,8 @@ type rebuilder struct {
 	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
 	keyIO *keyio.Handle
 
-	curKey          keyio.KeyAndTree
-	queue           []keyio.KeyAndTree
+	curKey          keyAndTree
+	queue           []keyAndTree
 	pendingAugments map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int
 }
 
@@ -179,7 +191,7 @@ func (o *rebuilder) rebuild(ctx context.Context) error {
 }
 
 func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key) {
-	o.queue = append(o.queue, keyio.KeyAndTree{
+	o.queue = append(o.queue, keyAndTree{
 		TreeID: tree,
 		Key:    key,
 	})
-- 
cgit v1.2.3-2-g168b


From 5fcb6f7ab77721faf0f89e00bf3087baba2d566c Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 19:41:39 -0700
Subject: rebuildnodes: Ignore the LOG_TREE for now

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 00dbad5..076710e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -117,7 +117,7 @@ func (o *rebuilder) rebuild(ctx context.Context) error {
 	o.pendingAugments = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
 	o.rebuilt.AddTree(ctx, btrfsprim.ROOT_TREE_OBJECTID)
 	o.rebuilt.AddTree(ctx, btrfsprim.CHUNK_TREE_OBJECTID)
-	o.rebuilt.AddTree(ctx, btrfsprim.TREE_LOG_OBJECTID)
+	//o.rebuilt.AddTree(ctx, btrfsprim.TREE_LOG_OBJECTID) // TODO(lukeshu): Special LOG_TREE handling
 	o.rebuilt.AddTree(ctx, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
 	for {
 		// Handle items in the queue
-- 
cgit v1.2.3-2-g168b


From b3264cfa1a48998bc9406c18339801c21a468f05 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 19:54:56 -0700
Subject: rebuildnodes: Less noisy logging about csums

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 076710e..1daa83d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -510,11 +510,13 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 //
 // interval is [beg, end)
 func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
-	if !o.rebuilt.AddTree(ctx, btrfsprim.CSUM_TREE_OBJECTID) {
+	treeID := btrfsprim.CSUM_TREE_OBJECTID
+	if !o.rebuilt.AddTree(ctx, treeID) {
 		o.queue = append(o.queue, o.curKey)
 		return
 	}
 
+	last := beg
 	for beg < end {
 		// Figure out which key we want.
 
@@ -531,20 +533,20 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 
 		})
 		if rbNode == nil {
-			o.wantAugment(ctx, btrfsprim.CSUM_TREE_OBJECTID, nil) // log an error
 			beg += btrfssum.BlockSize
 			continue
+		} else if last < beg {
+			dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
 		}
 		run := rbNode.Value.Sums
 		key := rbNode.Value.Key
-		treeID := btrfsprim.CSUM_TREE_OBJECTID
 
 		// Check if we already have it.
 
 		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
 		if _, ok := o.rebuilt.Search(ctx, treeID, key.Cmp); !ok {
 			// We need to insert it.
-			itemPtr, ok := o.rebuilt.Keys(btrfsprim.CSUM_TREE_OBJECTID).Load(key)
+			itemPtr, ok := o.rebuilt.Keys(treeID).Load(key)
 			if !ok {
 				// This is a panic because if we found it in `o.csums` then it has
 				// to be in some Node, and if we didn't find it from
@@ -555,6 +557,10 @@ func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr)
 		}
 
 		beg = run.Addr.Add(run.Size())
+		last = beg
+	}
+	if last < beg {
+		dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
 	}
 }
 
-- 
cgit v1.2.3-2-g168b


From ce1ea3ba56fb7738b4ca567930e357473991e7dd Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 23 Dec 2022 22:32:52 -0700
Subject: rebuildnodes: Rework the wantCsum and wantFileExt algorithms

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 348 ++++++++++-----------
 1 file changed, 169 insertions(+), 179 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 1daa83d..5badc33 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -15,11 +15,9 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"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/btrfssum"
 	"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/btrfsprogs/btrfsinspect"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildmappings"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
@@ -45,7 +43,6 @@ type rebuilder struct {
 	rebuilt *btrees.RebuiltTrees
 
 	graph graph.Graph
-	csums containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum]
 	keyIO *keyio.Handle
 
 	curKey          keyAndTree
@@ -65,19 +62,12 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		return nil, err
 	}
 
-	dlog.Info(ctx, "Indexing checksums...")
-	csums, _ := rebuildmappings.ExtractLogicalSums(ctx, nodeScanResults)
-	if csums == nil {
-		csums = new(containers.RBTree[containers.NativeOrdered[btrfsvol.LogicalAddr], btrfsinspect.SysExtentCSum])
-	}
-
 	dlog.Info(ctx, "Rebuilding node tree...")
 	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 	o := &rebuilder{
 		sb: *sb,
 
 		graph: nodeGraph,
-		csums: *csums,
 		keyIO: keyIO,
 	}
 	o.rebuilt = btrees.NewRebuiltTrees(*sb, nodeGraph, keyIO,
@@ -506,219 +496,219 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 	o.wantAugment(ctx, treeID, wants)
 }
 
-// func implements rebuildCallbacks.
-//
-// interval is [beg, end)
-func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
-	treeID := btrfsprim.CSUM_TREE_OBJECTID
+func (o *rebuilder) _wantRange(
+	ctx context.Context,
+	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
+	sizeFn func(btrfsprim.Key) (uint64, error),
+	beg, end uint64,
+) {
 	if !o.rebuilt.AddTree(ctx, treeID) {
 		o.queue = append(o.queue, o.curKey)
 		return
 	}
 
-	last := beg
-	for beg < end {
-		// Figure out which key we want.
-
-		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("csum for laddr=%v", beg))
-		rbNode := o.csums.Search(func(item btrfsinspect.SysExtentCSum) int {
-			switch {
-			case item.Sums.Addr > beg:
-				return -1
-			case item.Sums.Addr.Add(item.Sums.Size()) <= beg:
-				return 1
-			default:
-				return 0
-			}
-
-		})
-		if rbNode == nil {
-			beg += btrfssum.BlockSize
-			continue
-		} else if last < beg {
-			dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
-		}
-		run := rbNode.Value.Sums
-		key := rbNode.Value.Key
-
-		// Check if we already have it.
-
-		// .Search is more efficient than .Load, because it doesn't load the body (and we don't need the body).
-		if _, ok := o.rebuilt.Search(ctx, treeID, key.Cmp); !ok {
-			// We need to insert it.
-			itemPtr, ok := o.rebuilt.Keys(treeID).Load(key)
-			if !ok {
-				// This is a panic because if we found it in `o.csums` then it has
-				// to be in some Node, and if we didn't find it from
-				// btrfs.LookupCSum(), then that Node must be an orphan.
-				panic(fmt.Errorf("should not happen: no orphan contains %v", key))
-			}
-			o.wantAugment(ctx, treeID, o.rebuilt.LeafToRoots(ctx, treeID, itemPtr.Node))
-		}
-
-		beg = run.Addr.Add(run.Size())
-		last = beg
-	}
-	if last < beg {
-		dlog.Errorf(ctx, "augment(tree=%v): could not find csum items for %v-%v", treeID, last, beg)
-	}
-}
-
-// wantFileExt implements rebuildCallbacks.
-func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
-	if !o.rebuilt.AddTree(ctx, treeID) {
-		o.queue = append(o.queue, o.curKey)
-		return
-	}
-
-	min := btrfsprim.Key{
-		ObjectID: ino,
-		ItemType: btrfsitem.EXTENT_DATA_KEY,
-		Offset:   0,
+	// Step 1: Build a listing of the runs that we do have.
+	runMin := btrfsprim.Key{
+		ObjectID: objID,
+		ItemType: typ,
+		Offset:   0, // *NOT* `beg`
 	}
-	max := btrfsprim.Key{
-		ObjectID: ino,
-		ItemType: btrfsitem.EXTENT_DATA_KEY,
-		Offset:   uint64(size - 1),
+	runMax := btrfsprim.Key{
+		ObjectID: objID,
+		ItemType: typ,
+		Offset:   end - 1,
 	}
-	extKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
+	runKeys := o.rebuilt.SearchAll(ctx, treeID, func(key btrfsprim.Key) int {
 		switch {
-		case min.Cmp(key) < 0:
+		case runMin.Cmp(key) < 0:
 			return 1
-		case max.Cmp(key) > 0:
+		case runMax.Cmp(key) > 0:
 			return -1
 		default:
 			return 0
 		}
 	})
 
+	// Step 2: Build a listing of the gaps.
+	//
+	// Start with a gap of the whole range, then subtract each run
+	// from it.
 	type gap struct {
 		// range is [Beg,End)
-		Beg, End int64
+		Beg, End uint64
 	}
-	gaps := &containers.RBTree[containers.NativeOrdered[int64], gap]{
-		KeyFn: func(gap gap) containers.NativeOrdered[int64] {
-			return containers.NativeOrdered[int64]{Val: gap.Beg}
+	gaps := &containers.RBTree[containers.NativeOrdered[uint64], gap]{
+		KeyFn: func(gap gap) containers.NativeOrdered[uint64] {
+			return containers.NativeOrdered[uint64]{Val: gap.Beg}
 		},
 	}
 	gaps.Insert(gap{
-		Beg: 0,
-		End: size,
+		Beg: beg,
+		End: end,
 	})
-	for _, extKey := range extKeys {
-		extBody, ok := o.rebuilt.Load(ctx, treeID, extKey)
-		if !ok {
-			o.ioErr(ctx, fmt.Errorf("could not look up already-inserted item: tree=%v key=%v",
-				treeID, extKey))
+	for _, runKey := range runKeys {
+		runSize, err := sizeFn(runKey)
+		if err != nil {
+			o.fsErr(ctx, fmt.Errorf("tree=%v key=%v: %w", treeID, runKey, err))
+		}
+		if runSize == 0 {
+			continue
+		}
+		runBeg := runKey.Offset
+		runEnd := runBeg + runSize
+		if runEnd <= beg {
+			continue
 		}
-		switch extBody := extBody.(type) {
-		case btrfsitem.FileExtent:
-			extBeg := int64(extKey.Offset)
-			extSize, err := extBody.Size()
-			if err != nil {
-				o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, extKey, err))
-				continue
+		overlappingGaps := gaps.SearchRange(func(gap gap) int {
+			switch {
+			case gap.End <= runBeg:
+				return 1
+			case runEnd <= gap.Beg:
+				return -1
+			default:
+				return 0
 			}
-			extEnd := extBeg + extSize
-			overlappingGaps := gaps.SearchRange(func(gap gap) int {
-				switch {
-				case gap.End <= extBeg:
-					return 1
-				case extEnd <= gap.Beg:
-					return -1
-				default:
-					return 0
-				}
+		})
+		if len(overlappingGaps) == 0 {
+			continue
+		}
+		gapsBeg := overlappingGaps[0].Beg
+		gapsEnd := overlappingGaps[len(overlappingGaps)-1].End
+		for _, gap := range overlappingGaps {
+			gaps.Delete(containers.NativeOrdered[uint64]{Val: gap.Beg})
+		}
+		if gapsBeg < runBeg {
+			gaps.Insert(gap{
+				Beg: gapsBeg,
+				End: runBeg,
+			})
+		}
+		if gapsEnd > runEnd {
+			gaps.Insert(gap{
+				Beg: runEnd,
+				End: gapsEnd,
 			})
-			if len(overlappingGaps) == 0 {
-				continue
-			}
-			beg := overlappingGaps[0].Beg
-			end := overlappingGaps[len(overlappingGaps)-1].End
-			for _, gap := range overlappingGaps {
-				gaps.Delete(containers.NativeOrdered[int64]{Val: gap.Beg})
-			}
-			if beg < extBeg {
-				gaps.Insert(gap{
-					Beg: beg,
-					End: extBeg,
-				})
-			}
-			if end > extEnd {
-				gaps.Insert(gap{
-					Beg: extEnd,
-					End: end,
-				})
-			}
-		case btrfsitem.Error:
-			o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, extKey, extBody.Err))
-		default:
-			// This is a panic because the item decoder should not emit EXTENT_DATA
-			// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-			// this code also being updated.
-			panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", extBody))
 		}
 	}
+
+	// Step 2: Fill each gap.
 	_ = gaps.Walk(func(rbNode *containers.RBNode[gap]) error {
 		gap := rbNode.Value
-		min := btrfsprim.Key{
-			ObjectID: ino,
-			ItemType: btrfsitem.EXTENT_DATA_KEY,
-			Offset:   0,
+		last := gap.Beg
+		runMin := btrfsprim.Key{
+			ObjectID: objID,
+			ItemType: typ,
+			Offset:   0, // *NOT* `gap.Beg`
 		}
-		max := btrfsprim.Key{
-			ObjectID: ino,
-			ItemType: btrfsitem.EXTENT_DATA_KEY,
-			Offset:   uint64(gap.End - 1),
+		runMax := btrfsprim.Key{
+			ObjectID: objID,
+			ItemType: typ,
+			Offset:   gap.End - 1,
 		}
-		ctx := dlog.WithField(ctx, "want_key", fmt.Sprintf("file extent for tree=%v inode=%v bytes [%v, %v)", treeID, ino, gap.Beg, gap.End))
-		wants := make(containers.Set[btrfsvol.LogicalAddr])
 		o.rebuilt.Keys(treeID).Subrange(
-			func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+			func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 				switch {
-				case min.Cmp(k) < 0:
+				case runMin.Cmp(key) < 0:
 					return 1
-				case max.Cmp(k) > 0:
+				case runMax.Cmp(key) > 0:
 					return -1
 				default:
 					return 0
 				}
 			},
 			func(k btrfsprim.Key, v keyio.ItemPtr) bool {
-				itemBody, ok := o.keyIO.ReadItem(v)
-				if !ok {
-					o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", v))
+				runSize, err := sizeFn(k)
+				if err != nil {
+					o.fsErr(ctx, fmt.Errorf("tree=%v key=%v: %w", treeID, k, err))
+					return true
 				}
-				switch itemBody := itemBody.(type) {
-				case btrfsitem.FileExtent:
-					itemBeg := int64(k.Offset)
-					itemSize, err := itemBody.Size()
-					if err != nil {
-						o.fsErr(ctx, fmt.Errorf("FileExtent: tree=%v key=%v: %w", treeID, k, err))
-						break
-					}
-					itemEnd := itemBeg + itemSize
-					// We're being greedy and "wanting" any extent that has any overlap with
-					// the gap.  But maybe instead we sould only want extents that are
-					// *entirely* within the gap.  I'll have to run it on real filesystems
-					// to see what works better.
-					//
-					// TODO(lukeshu): Re-evaluate whether being greedy here is the right
-					// thing.
-					if itemEnd > gap.Beg && itemBeg < gap.End {
-						wants.InsertFrom(o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
-					}
-				case btrfsitem.Error:
-					o.fsErr(ctx, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, k, itemBody.Err))
-				default:
-					// This is a panic because the item decoder should not emit EXTENT_DATA
-					// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-					// this code also being updated.
-					panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", itemBody))
+				if runSize == 0 {
+					return true
+				}
+				runBeg := k.Offset
+				runEnd := runBeg + runSize
+				if runEnd <= gap.Beg {
+					return true
+				}
+
+				// TODO: This is dumb and greedy.
+				if last < runBeg {
+					// log an error
+					o.wantAugment(dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v, %v, %v-%v}",
+						treeID, objID, typ, last, runBeg)),
+						treeID, nil)
 				}
+				o.wantAugment(dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v, %v, %v-%v}",
+					treeID, objID, typ, gap.Beg, gap.End)),
+					treeID, o.rebuilt.LeafToRoots(ctx, treeID, v.Node))
+				last = runEnd
+
 				return true
 			})
-		o.wantAugment(ctx, treeID, wants)
+		if last < gap.End {
+			// log an error
+			o.wantAugment(dlog.WithField(ctx, "want_key", fmt.Sprintf("tree=%v key={%v, %v, %v-%v}",
+				treeID, objID, typ, last, gap.End)),
+				treeID, nil)
+		}
 		return nil
 	})
 }
+
+// func implements rebuildCallbacks.
+//
+// interval is [beg, end)
+func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
+	const treeID = btrfsprim.CSUM_TREE_OBJECTID
+	o._wantRange(ctx, treeID, btrfsprim.EXTENT_CSUM_OBJECTID, btrfsprim.EXTENT_CSUM_KEY,
+		func(key btrfsprim.Key) (uint64, error) {
+			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not load key: %v", key))
+			}
+			body, ok := o.keyIO.ReadItem(ptr)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not read item: %v", key))
+			}
+			switch body := body.(type) {
+			case btrfsitem.ExtentCSum:
+				return uint64(body.Size()), nil
+			case btrfsitem.Error:
+				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
+			default:
+				// This is a panic because the item decoder should not emit EXTENT_CSUM
+				// items as anything but btrfsitem.ExtentCSum or btrfsitem.Error without
+				// this code also being updated.
+				panic(fmt.Errorf("should not happen: EXTENT_CSUM item has unexpected type: %T", body))
+			}
+		},
+		uint64(beg), uint64(end))
+}
+
+// wantFileExt implements rebuildCallbacks.
+func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
+	o._wantRange(ctx, treeID, ino, btrfsprim.EXTENT_DATA_KEY,
+		func(key btrfsprim.Key) (uint64, error) {
+			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not load key: %v", key))
+			}
+			body, ok := o.keyIO.ReadItem(ptr)
+			if !ok {
+				panic(fmt.Errorf("should not happen: could not read item: %v", key))
+			}
+			switch body := body.(type) {
+			case btrfsitem.FileExtent:
+				size, err := body.Size()
+				return uint64(size), err
+			case btrfsitem.Error:
+				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
+			default:
+				// This is a panic because the item decoder should not emit EXTENT_DATA
+				// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
+				// this code also being updated.
+				panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", body))
+			}
+		},
+		0, uint64(size))
+}
-- 
cgit v1.2.3-2-g168b


From bc06b340b71a07a0bb500e1bb7e81ece19f1a9c5 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sat, 24 Dec 2022 04:38:24 -0700
Subject: rebuildnodes: Track item sizes at startup, to avoid extra i/o

---
 .../btrfsinspect/rebuildnodes/rebuild.go           | 59 +++++-----------------
 1 file changed, 13 insertions(+), 46 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 5badc33..66eaf1a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -51,7 +51,7 @@ type rebuilder struct {
 }
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	nodeGraph, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
@@ -63,7 +63,6 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Rebuilding node tree...")
-	keyIO := keyio.NewHandle(fs, *sb, nodeGraph)
 	o := &rebuilder{
 		sb: *sb,
 
@@ -499,7 +498,6 @@ func (o *rebuilder) wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 func (o *rebuilder) _wantRange(
 	ctx context.Context,
 	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
-	sizeFn func(btrfsprim.Key) (uint64, error),
 	beg, end uint64,
 ) {
 	if !o.rebuilt.AddTree(ctx, treeID) {
@@ -507,6 +505,18 @@ func (o *rebuilder) _wantRange(
 		return
 	}
 
+	sizeFn := func(key btrfsprim.Key) (uint64, error) {
+		ptr, ok := o.rebuilt.Keys(treeID).Load(key)
+		if !ok {
+			panic(fmt.Errorf("should not happen: could not load key: %v", key))
+		}
+		sizeAndErr, ok := o.keyIO.Sizes[ptr]
+		if !ok {
+			panic(fmt.Errorf("should not happen: %v item did not have a size recorded", typ))
+		}
+		return sizeAndErr.Size, sizeAndErr.Err
+	}
+
 	// Step 1: Build a listing of the runs that we do have.
 	runMin := btrfsprim.Key{
 		ObjectID: objID,
@@ -661,54 +671,11 @@ func (o *rebuilder) _wantRange(
 func (o *rebuilder) wantCSum(ctx context.Context, beg, end btrfsvol.LogicalAddr) {
 	const treeID = btrfsprim.CSUM_TREE_OBJECTID
 	o._wantRange(ctx, treeID, btrfsprim.EXTENT_CSUM_OBJECTID, btrfsprim.EXTENT_CSUM_KEY,
-		func(key btrfsprim.Key) (uint64, error) {
-			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not load key: %v", key))
-			}
-			body, ok := o.keyIO.ReadItem(ptr)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not read item: %v", key))
-			}
-			switch body := body.(type) {
-			case btrfsitem.ExtentCSum:
-				return uint64(body.Size()), nil
-			case btrfsitem.Error:
-				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
-			default:
-				// This is a panic because the item decoder should not emit EXTENT_CSUM
-				// items as anything but btrfsitem.ExtentCSum or btrfsitem.Error without
-				// this code also being updated.
-				panic(fmt.Errorf("should not happen: EXTENT_CSUM item has unexpected type: %T", body))
-			}
-		},
 		uint64(beg), uint64(end))
 }
 
 // wantFileExt implements rebuildCallbacks.
 func (o *rebuilder) wantFileExt(ctx context.Context, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
 	o._wantRange(ctx, treeID, ino, btrfsprim.EXTENT_DATA_KEY,
-		func(key btrfsprim.Key) (uint64, error) {
-			ptr, ok := o.rebuilt.Keys(treeID).Load(key)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not load key: %v", key))
-			}
-			body, ok := o.keyIO.ReadItem(ptr)
-			if !ok {
-				panic(fmt.Errorf("should not happen: could not read item: %v", key))
-			}
-			switch body := body.(type) {
-			case btrfsitem.FileExtent:
-				size, err := body.Size()
-				return uint64(size), err
-			case btrfsitem.Error:
-				return 0, fmt.Errorf("error decoding item: tree=%v key=%v: %w", treeID, key, body.Err)
-			default:
-				// This is a panic because the item decoder should not emit EXTENT_DATA
-				// items as anything but btrfsitem.FileExtent or btrfsitem.Error without
-				// this code also being updated.
-				panic(fmt.Errorf("should not happen: EXTENT_DATA item has unexpected type: %T", body))
-			}
-		},
 		0, uint64(size))
 }
-- 
cgit v1.2.3-2-g168b