From dcb9dbd93bc45cfce2ca8ca9eebf822a128caa97 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Mon, 2 Jan 2023 18:45:57 -0700
Subject: rebuildnodes/btrees: Cache failures to add a tree

---
 .../btrfsinspect/rebuildnodes/btrees/forrest.go           | 15 ++++++++++++---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go   |  9 +++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

(limited to 'lib')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 6bfa297..6f0ac96 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -105,9 +105,16 @@ func (ts *RebuiltForrest) Tree(ctx context.Context, treeID btrfsprim.ObjID) *Reb
 }
 
 func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, stack []btrfsprim.ObjID) (ok bool) {
-	if _, ok := ts.trees.Load(treeID); ok {
-		return true
+	if tree, ok := ts.trees.Load(treeID); ok {
+		return tree != nil
 	}
+	defer func() {
+		if !ok {
+			// Store a negative cache of this.  tree.AddRoot() for the ROOT or UUID
+			// trees will invalidate the negative cache.
+			ts.trees.Store(treeID, nil)
+		}
+	}()
 	if slices.Contains(treeID, stack) {
 		return false
 	}
@@ -173,7 +180,9 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 func (ts *RebuiltForrest) ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
 	ret := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr])
 	ts.trees.Range(func(treeID btrfsprim.ObjID, tree *RebuiltTree) bool {
-		ret[treeID] = tree.Roots
+		if tree != nil {
+			ret[treeID] = tree.Roots
+		}
 		return true
 	})
 	return ret
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index 678d25f..c9747ff 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -171,6 +171,15 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	stats.Leafs.N = len(leafToRoots)
 	progressWriter.Set(stats)
 	progressWriter.Done()
+
+	if (tree.ID == btrfsprim.ROOT_TREE_OBJECTID || tree.ID == btrfsprim.UUID_TREE_OBJECTID) && stats.AddedItems > 0 {
+		tree.forrest.trees.Range(func(otherTreeID btrfsprim.ObjID, otherTree *RebuiltTree) bool {
+			if otherTree == nil {
+				tree.forrest.trees.Delete(otherTreeID)
+			}
+			return true
+		})
+	}
 }
 
 // .Items() and .PotentialItems() //////////////////////////////////////////////////////////////////////////////////////
-- 
cgit v1.2.3-2-g168b