rebuildnodes: Handle the same node being added twice

2 files changed, 21 insertions, 11 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index cbe7df1..a4ccd5e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -34,6 +34,7 @@ type rebuiltTree struct {
 
 	// mutable
 	Roots containers.Set[btrfsvol.LogicalAddr]
+	Leafs containers.Set[btrfsvol.LogicalAddr]
 	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 }
 
@@ -175,10 +176,10 @@ func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, roo
 	}
 	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
 		progress(i)
-		roots := tree.leafToRoots[leaf]
-		if !roots.Has(rootNode) {
+		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
 			continue
 		}
+		tree.Leafs.Insert(leaf)
 		for j, item := range ts.keyIO.ReadNode(leaf).Data.BodyLeaf {
 			newPtr := keyio.ItemPtr{
 				Node: leaf,
@@ -234,6 +235,7 @@ func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, sta
 	tree := &rebuiltTree{
 		ID:    treeID,
 		Roots: make(containers.Set[btrfsvol.LogicalAddr]),
+		Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
 	}
 	var root btrfsvol.LogicalAddr
 	switch treeID {
@@ -404,11 +406,17 @@ func (ts *RebuiltTrees) LeafToRoots(ctx context.Context, treeID btrfsprim.ObjID,
 	if !ts.AddTree(ctx, treeID) {
 		return nil
 	}
+	if ts.graph.Nodes[leaf].Level != 0 {
+		panic(fmt.Errorf("should not happen: NodeToRoots(tree=%v, leaf=%v): not a leaf",
+			treeID, leaf))
+	}
 	ret := make(containers.Set[btrfsvol.LogicalAddr])
 	for root := range ts.trees[treeID].leafToRoots[leaf] {
-		if !ts.trees[treeID].Roots.Has(root) {
-			ret.Insert(root)
+		if ts.trees[treeID].Roots.Has(root) {
+			panic(fmt.Errorf("should not happen: NodeToRoots(tree=%v, leaf=%v): tree contains root=%v but not leaf",
+				treeID, leaf, root))
 		}
+		ret.Insert(root)
 	}
 	if len(ret) == 0 {
 		return nil
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)
 }