From 1f8734f894ac5fdaee25bba5b1e3ffd10b31ef4e Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 30 Dec 2022 19:45:07 -0700
Subject: rebuildnodes/btrees: Refactor to split the forrest from the trees

---
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 299 +++++++++++++++++++++
 1 file changed, 299 insertions(+)
 create mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
new file mode 100644
index 0000000..1fc1f52
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -0,0 +1,299 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrees
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/datawire/dlib/dlog"
+
+	"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/btrfsvol"
+	pkggraph "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/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+type RebuiltTree struct {
+	// static
+	ID        btrfsprim.ObjID
+	UUID      btrfsprim.UUID
+	Parent    *RebuiltTree
+	ParentGen btrfsprim.Generation // offset of this tree's root item
+	forrest   *RebuiltForrest
+
+	// all leafs (lvl=0) that pass .isOwnerOK, even if not in the tree
+	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
+	keys        containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
+
+	// mutable
+	Roots containers.Set[btrfsvol.LogicalAddr]
+	Leafs containers.Set[btrfsvol.LogicalAddr]
+	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
+}
+
+// initializaton (called by `RebuiltForrest.Tree()`) ///////////////////////////////////////////////////////////////////
+
+func (tree *RebuiltTree) indexLeafs(ctx context.Context) {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.substep", "index-nodes")
+
+	nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+
+	var stats textui.Portion[int]
+	stats.D = len(tree.forrest.graph.Nodes)
+	progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progress := func() {
+		stats.N = len(nodeToRoots)
+		progressWriter.Set(stats)
+	}
+
+	progress()
+	for _, node := range maps.SortedKeys(tree.forrest.graph.Nodes) {
+		tree.indexNode(ctx, node, nodeToRoots, progress, nil)
+	}
+	progressWriter.Done()
+
+	tree.leafToRoots = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+	for node, roots := range nodeToRoots {
+		if tree.forrest.graph.Nodes[node].Level == 0 && len(roots) > 0 {
+			tree.leafToRoots[node] = roots
+		}
+	}
+}
+
+func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
+	defer progress()
+	if err := ctx.Err(); err != nil {
+		return
+	}
+	if _, done := index[node]; done {
+		return
+	}
+	if slices.Contains(node, stack) {
+		// This is a panic because tree.forrest.graph.FinalCheck() should
+		// have already checked for loops.
+		panic("loop")
+	}
+	if !tree.isOwnerOK(tree.forrest.graph.Nodes[node].Owner, tree.forrest.graph.Nodes[node].Generation) {
+		index[node] = nil
+		return
+	}
+
+	// tree.leafToRoots
+	stack = append(stack, node)
+	var roots containers.Set[btrfsvol.LogicalAddr]
+	kps := slices.RemoveAllFunc(tree.forrest.graph.EdgesTo[node], func(kp *pkggraph.Edge) bool {
+		return !tree.isOwnerOK(tree.forrest.graph.Nodes[kp.FromNode].Owner, tree.forrest.graph.Nodes[kp.FromNode].Generation)
+	})
+	for _, kp := range kps {
+		tree.indexNode(ctx, kp.FromNode, index, progress, stack)
+		if len(index[kp.FromNode]) > 0 {
+			if roots == nil {
+				roots = make(containers.Set[btrfsvol.LogicalAddr])
+			}
+			roots.InsertFrom(index[kp.FromNode])
+		}
+	}
+	if roots == nil {
+		roots = containers.NewSet[btrfsvol.LogicalAddr](node)
+	}
+	index[node] = roots
+
+	// tree.keys
+	for i, key := range tree.forrest.graph.Nodes[node].Items {
+		if oldPtr, ok := tree.keys.Load(key); !ok || tree.shouldReplace(oldPtr.Node, node) {
+			tree.keys.Store(key, keyio.ItemPtr{
+				Node: node,
+				Idx:  i,
+			})
+		}
+	}
+}
+
+// isOwnerOK returns whether it is permissible for a node with
+// .Head.Owner=owner to be in this tree.
+func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generation) bool {
+	for {
+		if owner == tree.ID {
+			return true
+		}
+		if tree.Parent == nil || gen >= tree.ParentGen {
+			return false
+		}
+		tree = tree.Parent
+	}
+}
+
+// .AddRoot() //////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+type rootStats struct {
+	Leafs         textui.Portion[int]
+	AddedItems    int
+	ReplacedItems int
+}
+
+func (s rootStats) String() string {
+	return textui.Sprintf("%v (added %v items, replaced %v items)",
+		s.Leafs, s.AddedItems, s.ReplacedItems)
+}
+
+// AddRoot adds an additional root node to the tree.  It is useful to
+// call .AddRoot() to re-attach part of the tree that has been broken
+// off.
+func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
+	tree.Roots.Insert(rootNode)
+
+	var stats rootStats
+	stats.Leafs.D = len(tree.leafToRoots)
+	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
+		stats.Leafs.N = i
+		progressWriter.Set(stats)
+		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
+			continue
+		}
+		tree.Leafs.Insert(leaf)
+		for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+			newPtr := keyio.ItemPtr{
+				Node: leaf,
+				Idx:  j,
+			}
+			if oldPtr, exists := tree.Items.Load(itemKey); !exists {
+				tree.Items.Store(itemKey, newPtr)
+				stats.AddedItems++
+			} else if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
+				tree.Items.Store(itemKey, newPtr)
+				stats.ReplacedItems++
+			}
+			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
+			progressWriter.Set(stats)
+		}
+	}
+	stats.Leafs.N = len(tree.leafToRoots)
+	progressWriter.Set(stats)
+	progressWriter.Done()
+}
+
+func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bool {
+	oldDist, _ := tree.COWDistance(tree.forrest.graph.Nodes[oldNode].Owner)
+	newDist, _ := tree.COWDistance(tree.forrest.graph.Nodes[newNode].Owner)
+	switch {
+	case newDist < oldDist:
+		// Replace the old one with the new lower-dist one.
+		return true
+	case newDist > oldDist:
+		// Retain the old lower-dist one.
+		return false
+	default:
+		oldGen := tree.forrest.graph.Nodes[oldNode].Generation
+		newGen := tree.forrest.graph.Nodes[newNode].Generation
+		switch {
+		case newGen > oldGen:
+			// Replace the old one with the new higher-gen one.
+			return true
+		case newGen < oldGen:
+			// Retain the old higher-gen one.
+			return false
+		default:
+			// This is a panic because I'm not really sure what the best way to
+			// handle this is, and so if this happens I want the program to crash
+			// and force me to figure out how to handle it.
+			panic(fmt.Errorf("dup nodes in tree=%v: old=%v=%v ; new=%v=%v",
+				tree.ID,
+				oldNode, tree.forrest.graph.Nodes[oldNode],
+				newNode, tree.forrest.graph.Nodes[newNode]))
+		}
+	}
+}
+
+// main public API /////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// COWDistance returns how many COW-snapshots down the 'tree' is from
+// the 'parent'.
+func (tree *RebuiltTree) COWDistance(parentID btrfsprim.ObjID) (dist int, ok bool) {
+	for {
+		if parentID == tree.ID {
+			return dist, true
+		}
+		if tree.Parent == nil {
+			return 0, false
+		}
+		tree = tree.Parent
+		dist++
+	}
+}
+
+// Resolve a key to a keyio.ItemPtr.
+func (tree *RebuiltTree) Resolve(key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
+	return tree.Items.Load(key)
+}
+
+// Load reads an item from a tree.
+func (tree *RebuiltTree) Load(ctx context.Context, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
+	ptr, ok := tree.Resolve(key)
+	if !ok {
+		return nil, false
+	}
+	return tree.forrest.keyIO.ReadItem(ctx, ptr)
+}
+
+// Search searches for an item from a tree.
+func (tree *RebuiltTree) Search(fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
+	k, _, ok := tree.Items.Search(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+		return fn(k)
+	})
+	return k, ok
+}
+
+// Search searches for a range of items from a tree.
+func (tree *RebuiltTree) SearchAll(fn func(btrfsprim.Key) int) []btrfsprim.Key {
+	kvs := tree.Items.SearchAll(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+		return fn(k)
+	})
+	if len(kvs) == 0 {
+		return nil
+	}
+	ret := make([]btrfsprim.Key, len(kvs))
+	for i := range kvs {
+		ret[i] = kvs[i].K
+	}
+	return ret
+}
+
+// LeafToRoots returns the list of potential roots (to pass to
+// .AddRoot) that include a given leaf-node.
+func (tree *RebuiltTree) LeafToRoots(leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
+	if tree.forrest.graph.Nodes[leaf].Level != 0 {
+		panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): not a leaf",
+			tree.ID, leaf))
+	}
+	ret := make(containers.Set[btrfsvol.LogicalAddr])
+	for root := range tree.leafToRoots[leaf] {
+		if tree.Roots.Has(root) {
+			panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): tree contains root=%v but not leaf",
+				tree.ID, leaf, root))
+		}
+		ret.Insert(root)
+	}
+	if len(ret) == 0 {
+		return nil
+	}
+	return ret
+}
+
+// Keys returns a map of all keys in node that would be valid in this tree.
+//
+// Do not mutate the returned map; it is a pointer to the
+// RebuiltTree's internal map!
+func (tree *RebuiltTree) Keys() *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	return &tree.keys
+}
-- 
cgit v1.2.3-54-g00ecf


From e18c0e92ba35bb863f7375b190b0448d5fa65d33 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 30 Dec 2022 21:52:57 -0700
Subject: rebuildnodes/btrees: Allow item rbtrees to be evicted

---
 .../btrfsinspect/rebuildnodes/btrees/forrest.go    |   9 +-
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 148 +++++++++++++++------
 .../btrfsinspect/rebuildnodes/rebuild.go           |  20 +--
 lib/containers/sortedmap.go                        |   4 +
 lib/textui/log.go                                  |  16 ++-
 5 files changed, 135 insertions(+), 62 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 791e646..c0b7698 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -17,6 +17,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 // RebuiltForrest is an abstraction for rebuilding and accessing
@@ -61,7 +62,9 @@ type RebuiltForrest struct {
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
 
 	// mutable
-	trees map[btrfsprim.ObjID]*RebuiltTree
+	trees    map[btrfsprim.ObjID]*RebuiltTree
+	allItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
+	incItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 }
 
 // NewRebuiltForrest returns a new RebuiltForrest instance.  All of
@@ -81,7 +84,9 @@ func NewRebuiltForrest(
 		cbLookupRoot: cbLookupRoot,
 		cbLookupUUID: cbLookupUUID,
 
-		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
+		trees:    make(map[btrfsprim.ObjID]*RebuiltTree),
+		allItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(16)),
+		incItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(16)),
 	}
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index 1fc1f52..acc71db 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -30,14 +30,12 @@ type RebuiltTree struct {
 	ParentGen btrfsprim.Generation // offset of this tree's root item
 	forrest   *RebuiltForrest
 
-	// all leafs (lvl=0) that pass .isOwnerOK, even if not in the tree
+	// all leafs (lvl=0) that pass .isOwnerOK, whether or not they're  in the tree
 	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
-	keys        containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 
 	// mutable
 	Roots containers.Set[btrfsvol.LogicalAddr]
 	Leafs containers.Set[btrfsvol.LogicalAddr]
-	Items containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 }
 
 // initializaton (called by `RebuiltForrest.Tree()`) ///////////////////////////////////////////////////////////////////
@@ -106,16 +104,6 @@ func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAdd
 		roots = containers.NewSet[btrfsvol.LogicalAddr](node)
 	}
 	index[node] = roots
-
-	// tree.keys
-	for i, key := range tree.forrest.graph.Nodes[node].Items {
-		if oldPtr, ok := tree.keys.Load(key); !ok || tree.shouldReplace(oldPtr.Node, node) {
-			tree.keys.Store(key, keyio.ItemPtr{
-				Node: node,
-				Idx:  i,
-			})
-		}
-	}
 }
 
 // isOwnerOK returns whether it is permissible for a node with
@@ -135,14 +123,14 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 // .AddRoot() //////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 type rootStats struct {
-	Leafs         textui.Portion[int]
-	AddedItems    int
-	ReplacedItems int
+	Leafs      textui.Portion[int]
+	AddedLeafs int
+	AddedItems int
 }
 
 func (s rootStats) String() string {
-	return textui.Sprintf("%v (added %v items, replaced %v items)",
-		s.Leafs, s.AddedItems, s.ReplacedItems)
+	return textui.Sprintf("%v (added %v leafs, added %v items)",
+		s.Leafs, s.AddedLeafs, s.AddedItems)
 }
 
 // AddRoot adds an additional root node to the tree.  It is useful to
@@ -158,29 +146,109 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
 		stats.Leafs.N = i
 		progressWriter.Set(stats)
+
 		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
 			continue
 		}
+
 		tree.Leafs.Insert(leaf)
+		stats.AddedLeafs++
+		progressWriter.Set(stats)
+
+		for _, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
+			stats.AddedItems++
+			progressWriter.Set(stats)
+		}
+	}
+	stats.Leafs.N = len(tree.leafToRoots)
+	progressWriter.Set(stats)
+	progressWriter.Done()
+}
+
+// .Items() and .PotentialItems() //////////////////////////////////////////////////////////////////////////////////////
+
+// Items returns a map of the items contained in this tree.
+//
+// Do not mutate the returned map; it is a pointer to the
+// RebuiltTree's internal map!
+func (tree *RebuiltTree) Items(ctx context.Context) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-inc-items", fmt.Sprintf("tree=%v", tree.ID))
+	return tree.items(ctx, tree.forrest.incItems, maps.SortedKeys(tree.Leafs))
+}
+
+// PotentialItems returns a map of items that could be added to this
+// tree with .AddRoot().
+//
+// Do not mutate the returned map; it is a pointer to the
+// RebuiltTree's internal map!
+func (tree *RebuiltTree) PotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-all-items", fmt.Sprintf("tree=%v", tree.ID))
+	return tree.items(ctx, tree.forrest.allItems, maps.SortedKeys(tree.leafToRoots))
+}
+
+type itemIndex struct {
+	NumDups int
+	Leafs   containers.Set[btrfsvol.LogicalAddr]
+	Items   containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
+}
+
+type itemStats struct {
+	Leafs    textui.Portion[int]
+	NumItems int
+	NumDups  int
+}
+
+func (s itemStats) String() string {
+	return textui.Sprintf("%v (%v items, %v dups)",
+		s.Leafs, s.NumItems, s.NumDups)
+}
+
+func (tree *RebuiltTree) items(ctx context.Context, cache *containers.LRUCache[btrfsprim.ObjID, *itemIndex], leafs []btrfsvol.LogicalAddr) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	index := cache.GetOrElse(tree.ID, func() *itemIndex {
+		return &itemIndex{
+			Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
+		}
+	})
+
+	var stats itemStats
+	stats.Leafs.D = len(leafs)
+	progressWriter := textui.NewProgress[itemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progress := func(doneLeafs int) {
+		stats.Leafs.N = doneLeafs
+		stats.NumItems = index.Items.Len()
+		stats.NumDups = index.NumDups
+		progressWriter.Set(stats)
+	}
+
+	for i, leaf := range leafs {
+		if index.Leafs.Has(leaf) {
+			continue
+		}
+		progress(i)
+		index.Leafs.Insert(leaf)
 		for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
 			newPtr := keyio.ItemPtr{
 				Node: leaf,
 				Idx:  j,
 			}
-			if oldPtr, exists := tree.Items.Load(itemKey); !exists {
-				tree.Items.Store(itemKey, newPtr)
-				stats.AddedItems++
-			} else if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
-				tree.Items.Store(itemKey, newPtr)
-				stats.ReplacedItems++
+			if oldPtr, exists := index.Items.Load(itemKey); !exists {
+				index.Items.Store(itemKey, newPtr)
+			} else {
+				index.NumDups++
+				if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
+					index.Items.Store(itemKey, newPtr)
+				}
 			}
-			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
-			progressWriter.Set(stats)
+			progress(i)
 		}
 	}
-	stats.Leafs.N = len(tree.leafToRoots)
-	progressWriter.Set(stats)
-	progressWriter.Done()
+	if stats.Leafs.N > 0 {
+		progress(len(leafs))
+		progressWriter.Done()
+	}
+
+	return &index.Items
 }
 
 func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bool {
@@ -233,13 +301,13 @@ func (tree *RebuiltTree) COWDistance(parentID btrfsprim.ObjID) (dist int, ok boo
 }
 
 // Resolve a key to a keyio.ItemPtr.
-func (tree *RebuiltTree) Resolve(key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
-	return tree.Items.Load(key)
+func (tree *RebuiltTree) Resolve(ctx context.Context, key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
+	return tree.Items(ctx).Load(key)
 }
 
 // Load reads an item from a tree.
 func (tree *RebuiltTree) Load(ctx context.Context, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
-	ptr, ok := tree.Resolve(key)
+	ptr, ok := tree.Resolve(ctx, key)
 	if !ok {
 		return nil, false
 	}
@@ -247,16 +315,16 @@ func (tree *RebuiltTree) Load(ctx context.Context, key btrfsprim.Key) (item btrf
 }
 
 // Search searches for an item from a tree.
-func (tree *RebuiltTree) Search(fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
-	k, _, ok := tree.Items.Search(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+func (tree *RebuiltTree) Search(ctx context.Context, fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
+	k, _, ok := tree.Items(ctx).Search(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
 		return fn(k)
 	})
 	return k, ok
 }
 
 // Search searches for a range of items from a tree.
-func (tree *RebuiltTree) SearchAll(fn func(btrfsprim.Key) int) []btrfsprim.Key {
-	kvs := tree.Items.SearchAll(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+func (tree *RebuiltTree) SearchAll(ctx context.Context, fn func(btrfsprim.Key) int) []btrfsprim.Key {
+	kvs := tree.Items(ctx).SearchAll(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
 		return fn(k)
 	})
 	if len(kvs) == 0 {
@@ -289,11 +357,3 @@ func (tree *RebuiltTree) LeafToRoots(leaf btrfsvol.LogicalAddr) containers.Set[b
 	}
 	return ret
 }
-
-// Keys returns a map of all keys in node that would be valid in this tree.
-//
-// Do not mutate the returned map; it is a pointer to the
-// RebuiltTree's internal map!
-func (tree *RebuiltTree) Keys() *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
-	return &tree.keys
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index b4feacf..85270d7 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -413,7 +413,7 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	if key, ok := o.rebuilt.Tree(ctx, treeID).Search(func(key btrfsprim.Key) int {
+	if key, ok := o.rebuilt.Tree(ctx, treeID).Search(ctx, func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	}); ok {
@@ -423,7 +423,7 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 	// OK, we need to insert it
 
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Tree(ctx, treeID).Keys().Subrange(
+	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).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.Tree(ctx, treeID).LeafToRoots(v.Node))
@@ -453,14 +453,14 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, tgt bt
 
 	// check if we already have it
 
-	if _, ok := o.rebuilt.Tree(ctx, treeID).Search(tgt.Cmp); ok {
+	if _, ok := o.rebuilt.Tree(ctx, treeID).Search(ctx, tgt.Cmp); ok {
 		return true
 	}
 
 	// OK, we need to insert it
 
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Tree(ctx, treeID).Keys().Subrange(
+	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { return tgt.Cmp(k) },
 		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
 			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(v.Node))
@@ -482,12 +482,12 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	keys := o.rebuilt.Tree(ctx, treeID).SearchAll(func(key btrfsprim.Key) int {
+	keys := o.rebuilt.Tree(ctx, treeID).SearchAll(ctx, func(key btrfsprim.Key) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	})
 	for _, itemKey := range keys {
-		itemPtr, ok := o.rebuilt.Tree(ctx, treeID).Resolve(itemKey)
+		itemPtr, ok := o.rebuilt.Tree(ctx, treeID).Resolve(ctx, itemKey)
 		if !ok {
 			o.ioErr(ctx, fmt.Errorf("could not resolve previously read item: %v", itemKey))
 		}
@@ -499,7 +499,7 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 	// OK, we need to insert it
 
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Tree(ctx, treeID).Keys().Subrange(
+	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
 		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 			if fn(v) {
@@ -536,7 +536,7 @@ func (o *rebuilder) _wantRange(
 	}
 
 	sizeFn := func(key btrfsprim.Key) (uint64, error) {
-		ptr, ok := o.rebuilt.Tree(ctx, treeID).Keys().Load(key)
+		ptr, ok := o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Load(key)
 		if !ok {
 			panic(fmt.Errorf("should not happen: could not load key: %v", key))
 		}
@@ -558,7 +558,7 @@ func (o *rebuilder) _wantRange(
 		ItemType: typ,
 		Offset:   end - 1,
 	}
-	runKeys := o.rebuilt.Tree(ctx, treeID).SearchAll(func(key btrfsprim.Key) int {
+	runKeys := o.rebuilt.Tree(ctx, treeID).SearchAll(ctx, func(key btrfsprim.Key) int {
 		switch {
 		case runMin.Cmp(key) < 0:
 			return 1
@@ -645,7 +645,7 @@ func (o *rebuilder) _wantRange(
 			ItemType: typ,
 			Offset:   gap.End - 1,
 		}
-		o.rebuilt.Tree(ctx, treeID).Keys().Subrange(
+		o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
 			func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 				switch {
 				case runMin.Cmp(key) < 0:
diff --git a/lib/containers/sortedmap.go b/lib/containers/sortedmap.go
index d6d2f9e..52308c9 100644
--- a/lib/containers/sortedmap.go
+++ b/lib/containers/sortedmap.go
@@ -92,3 +92,7 @@ func (m *SortedMap[K, V]) SearchAll(fn func(K, V) int) []OrderedKV[K, V] {
 		return fn(kv.K, kv.V)
 	})
 }
+
+func (m *SortedMap[K, V]) Len() int {
+	return m.inner.Len()
+}
diff --git a/lib/textui/log.go b/lib/textui/log.go
index 801e9eb..da303dd 100644
--- a/lib/textui/log.go
+++ b/lib/textui/log.go
@@ -317,18 +317,22 @@ func fieldOrd(key string) int {
 		return -25
 	// step=rebuild (any substep)
 	case "btrfsinspect.rebuild-nodes.rebuild.want.key":
-		return -7
+		return -9
 	case "btrfsinspect.rebuild-nodes.rebuild.want.reason":
-		return -6
+		return -8
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree":
-		return -5
+		return -7
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree.substep":
-		return -4
+		return -6
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.key":
-		return -3
+		return -5
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.reason":
-		return -2
+		return -4
 	case "btrfsinspect.rebuild-nodes.rebuild.add-root":
+		return -3
+	case "btrfsinspect.rebuild-nodes.rebuild.index-inc-items":
+		return -2
+	case "btrfsinspect.rebuild-nodes.rebuild.index-all-items":
 		return -1
 
 	// other ///////////////////////////////////////////////////////////////
-- 
cgit v1.2.3-54-g00ecf


From 3c3ec3e4ebb93b8e09a5a93bd26ace84f271223e Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Fri, 30 Dec 2022 22:36:05 -0700
Subject: rebuildnodes/btrees.RebuiltTree: Try to remove methods

Now that .Items() is public, some of the search methods are superfluous,
and in fact all .SearchAll calls would be more efficient as
.Items.Subrange calls.

And rename .Load to .ReadItem, so that grepping for it doesn't mix up with
.Items.Load.
---
 .../btrfsinspect/rebuildnodes/btrees/tree.go       |  34 +-----
 .../btrfsinspect/rebuildnodes/rebuild.go           | 133 +++++++++++----------
 2 files changed, 72 insertions(+), 95 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index acc71db..ffbaa0f 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -300,43 +300,15 @@ func (tree *RebuiltTree) COWDistance(parentID btrfsprim.ObjID) (dist int, ok boo
 	}
 }
 
-// Resolve a key to a keyio.ItemPtr.
-func (tree *RebuiltTree) Resolve(ctx context.Context, key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
-	return tree.Items(ctx).Load(key)
-}
-
-// Load reads an item from a tree.
-func (tree *RebuiltTree) Load(ctx context.Context, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
-	ptr, ok := tree.Resolve(ctx, key)
+// ReadItem reads an item from a tree.
+func (tree *RebuiltTree) ReadItem(ctx context.Context, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
+	ptr, ok := tree.Items(ctx).Load(key)
 	if !ok {
 		return nil, false
 	}
 	return tree.forrest.keyIO.ReadItem(ctx, ptr)
 }
 
-// Search searches for an item from a tree.
-func (tree *RebuiltTree) Search(ctx context.Context, fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
-	k, _, ok := tree.Items(ctx).Search(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
-		return fn(k)
-	})
-	return k, ok
-}
-
-// Search searches for a range of items from a tree.
-func (tree *RebuiltTree) SearchAll(ctx context.Context, fn func(btrfsprim.Key) int) []btrfsprim.Key {
-	kvs := tree.Items(ctx).SearchAll(func(k btrfsprim.Key, _ keyio.ItemPtr) int {
-		return fn(k)
-	})
-	if len(kvs) == 0 {
-		return nil
-	}
-	ret := make([]btrfsprim.Key, len(kvs))
-	for i := range kvs {
-		ret[i] = kvs[i].K
-	}
-	return ret
-}
-
 // LeafToRoots returns the list of potential roots (to pass to
 // .AddRoot) that include a given leaf-node.
 func (tree *RebuiltTree) LeafToRoots(leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 85270d7..dcf77b8 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -139,7 +139,7 @@ func (o *rebuilder) Rebuild(_ctx context.Context) error {
 				return err
 			}
 			o.curKey = key
-			itemBody, ok := o.rebuilt.Tree(itemCtx, key.TreeID).Load(itemCtx, key.Key)
+			itemBody, ok := o.rebuilt.Tree(itemCtx, key.TreeID).ReadItem(itemCtx, key.Key)
 			if !ok {
 				o.ioErr(itemCtx, fmt.Errorf("could not read previously read item: %v", key))
 			}
@@ -206,7 +206,7 @@ func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (off
 		o.itemQueue.Insert(o.curKey)
 		return 0, btrfsitem.Root{}, false
 	}
-	itemBody, ok := o.rebuilt.Tree(ctx, btrfsprim.ROOT_TREE_OBJECTID).Load(ctx, key)
+	itemBody, ok := o.rebuilt.Tree(ctx, btrfsprim.ROOT_TREE_OBJECTID).ReadItem(ctx, key)
 	if !ok {
 		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
 	}
@@ -231,7 +231,7 @@ func (o *rebuilder) cbLookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id b
 		o.itemQueue.Insert(o.curKey)
 		return 0, false
 	}
-	itemBody, ok := o.rebuilt.Tree(ctx, btrfsprim.UUID_TREE_OBJECTID).Load(ctx, key)
+	itemBody, ok := o.rebuilt.Tree(ctx, btrfsprim.UUID_TREE_OBJECTID).ReadItem(ctx, key)
 	if !ok {
 		o.ioErr(ctx, fmt.Errorf("could not read previously read item: %v", key))
 	}
@@ -413,7 +413,7 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	if key, ok := o.rebuilt.Tree(ctx, treeID).Search(ctx, func(key btrfsprim.Key) int {
+	if key, _, ok := o.rebuilt.Tree(ctx, treeID).Items(ctx).Search(func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 		key.Offset = 0
 		return tgt.Cmp(key)
 	}); ok {
@@ -453,7 +453,7 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, tgt bt
 
 	// check if we already have it
 
-	if _, ok := o.rebuilt.Tree(ctx, treeID).Search(ctx, tgt.Cmp); ok {
+	if _, ok := o.rebuilt.Tree(ctx, treeID).Items(ctx).Load(tgt); ok {
 		return true
 	}
 
@@ -482,18 +482,20 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 		ObjectID: objID,
 		ItemType: typ,
 	}
-	keys := o.rebuilt.Tree(ctx, treeID).SearchAll(ctx, func(key btrfsprim.Key) int {
-		key.Offset = 0
-		return tgt.Cmp(key)
-	})
-	for _, itemKey := range keys {
-		itemPtr, ok := o.rebuilt.Tree(ctx, treeID).Resolve(ctx, itemKey)
-		if !ok {
-			o.ioErr(ctx, fmt.Errorf("could not resolve previously read item: %v", itemKey))
-		}
-		if fn(itemPtr) {
-			return
-		}
+	found := false
+	o.rebuilt.Tree(ctx, treeID).Items(ctx).Subrange(
+		func(key btrfsprim.Key, _ keyio.ItemPtr) int {
+			key.Offset = 0
+			return tgt.Cmp(key)
+		},
+		func(_ btrfsprim.Key, itemPtr keyio.ItemPtr) bool {
+			if fn(itemPtr) {
+				found = true
+			}
+			return !found
+		})
+	if found {
+		return
 	}
 
 	// OK, we need to insert it
@@ -558,16 +560,6 @@ func (o *rebuilder) _wantRange(
 		ItemType: typ,
 		Offset:   end - 1,
 	}
-	runKeys := o.rebuilt.Tree(ctx, treeID).SearchAll(ctx, func(key btrfsprim.Key) int {
-		switch {
-		case runMin.Cmp(key) < 0:
-			return 1
-		case runMax.Cmp(key) > 0:
-			return -1
-		default:
-			return 0
-		}
-	})
 
 	// Step 2: Build a listing of the gaps.
 	//
@@ -586,50 +578,63 @@ func (o *rebuilder) _wantRange(
 		Beg: beg,
 		End: end,
 	})
-	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
-		}
-		overlappingGaps := gaps.SearchRange(func(gap gap) int {
+	o.rebuilt.Tree(ctx, treeID).Items(ctx).Subrange(
+		func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 			switch {
-			case gap.End <= runBeg:
+			case runMin.Cmp(key) < 0:
 				return 1
-			case runEnd <= gap.Beg:
+			case runMax.Cmp(key) > 0:
 				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,
+		},
+		func(runKey btrfsprim.Key, _ keyio.ItemPtr) bool {
+			runSize, err := sizeFn(runKey)
+			if err != nil {
+				o.fsErr(ctx, fmt.Errorf("tree=%v key=%v: %w", treeID, runKey, err))
+				return true
+			}
+			if runSize == 0 {
+				return true
+			}
+			runBeg := runKey.Offset
+			runEnd := runBeg + runSize
+			if runEnd <= beg {
+				return true
+			}
+			overlappingGaps := gaps.SearchRange(func(gap gap) int {
+				switch {
+				case gap.End <= runBeg:
+					return 1
+				case runEnd <= gap.Beg:
+					return -1
+				default:
+					return 0
+				}
 			})
-		}
-	}
+			if len(overlappingGaps) == 0 {
+				return true
+			}
+			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,
+				})
+			}
+			return true
+		})
 
 	// Step 2: Fill each gap.
 	_ = gaps.Walk(func(rbNode *containers.RBNode[gap]) error {
-- 
cgit v1.2.3-54-g00ecf


From e1302ac1f2685db057b7ecafe70f57ad17d533dc Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sat, 31 Dec 2022 12:45:16 -0700
Subject: rebuildnodes: Parallelize I/O and CPU

---
 .../btrfsinspect/rebuildnodes/btrees/forrest.go    | 23 ++++----
 .../btrfsinspect/rebuildnodes/btrees/tree.go       |  6 +++
 .../btrfsinspect/rebuildnodes/rebuild.go           | 63 ++++++++++++++--------
 3 files changed, 60 insertions(+), 32 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index ef6da6f..60f922d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -1,4 +1,4 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
@@ -62,7 +62,7 @@ type RebuiltForrest struct {
 	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
 
 	// mutable
-	trees    map[btrfsprim.ObjID]*RebuiltTree
+	trees    containers.SyncMap[btrfsprim.ObjID, *RebuiltTree]
 	allItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 	incItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 }
@@ -84,7 +84,6 @@ func NewRebuiltForrest(
 		cbLookupRoot: cbLookupRoot,
 		cbLookupUUID: cbLookupUUID,
 
-		trees:    make(map[btrfsprim.ObjID]*RebuiltTree),
 		allItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 		incItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 	}
@@ -99,11 +98,12 @@ func (ts *RebuiltForrest) Tree(ctx context.Context, treeID btrfsprim.ObjID) *Reb
 	if !ts.addTree(ctx, treeID, nil) {
 		return nil
 	}
-	return ts.trees[treeID]
+	tree, _ := ts.trees.Load(treeID)
+	return tree
 }
 
 func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, stack []btrfsprim.ObjID) (ok bool) {
-	if _, ok := ts.trees[treeID]; ok {
+	if _, ok := ts.trees.Load(treeID); ok {
 		return true
 	}
 	if slices.Contains(treeID, stack) {
@@ -151,12 +151,12 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			if !ts.addTree(ctx, parentID, append(stack, treeID)) {
 				return false
 			}
-			tree.Parent = ts.trees[parentID]
+			tree.Parent, _ = ts.trees.Load(parentID)
 		}
 	}
 	tree.indexLeafs(ctx)
 
-	ts.trees[treeID] = tree
+	ts.trees.Store(treeID, tree)
 	if root != 0 {
 		tree.AddRoot(ctx, root)
 	}
@@ -170,9 +170,10 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 // Do not mutate the set of roots for a tree; it is a pointer to the
 // RebuiltForrest's internal set!
 func (ts *RebuiltForrest) ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
-	ret := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], len(ts.trees))
-	for treeID := range ts.trees {
-		ret[treeID] = ts.trees[treeID].Roots
-	}
+	ret := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr])
+	ts.trees.Range(func(treeID btrfsprim.ObjID, tree *RebuiltTree) bool {
+		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 ffbaa0f..1daeefc 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -7,6 +7,7 @@ package btrees
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/datawire/dlib/dlog"
@@ -34,6 +35,7 @@ type RebuiltTree struct {
 	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
 
 	// mutable
+	mu    sync.Mutex
 	Roots containers.Set[btrfsvol.LogicalAddr]
 	Leafs containers.Set[btrfsvol.LogicalAddr]
 }
@@ -137,6 +139,8 @@ func (s rootStats) String() string {
 // call .AddRoot() to re-attach part of the tree that has been broken
 // off.
 func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
+	tree.mu.Lock()
+	defer tree.mu.Unlock()
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
 	tree.Roots.Insert(rootNode)
 
@@ -205,6 +209,8 @@ func (s itemStats) String() string {
 }
 
 func (tree *RebuiltTree) items(ctx context.Context, cache *containers.LRUCache[btrfsprim.ObjID, *itemIndex], leafs []btrfsvol.LogicalAddr) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+	tree.mu.Lock()
+	defer tree.mu.Unlock()
 	index := cache.GetOrElse(tree.ID, func() *itemIndex {
 		return &itemIndex{
 			Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index dcf77b8..cefa668 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -11,6 +11,7 @@ import (
 	"sort"
 	"time"
 
+	"github.com/datawire/dlib/dgroup"
 	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
@@ -130,30 +131,50 @@ func (o *rebuilder) Rebuild(_ctx context.Context) error {
 		progress.D = len(itemQueue)
 		progressWriter := textui.NewProgress[textui.Portion[int]](stepCtx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 		stepCtx = dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.substep.progress", &progress)
-		for i, key := range itemQueue {
-			itemCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.process.item", key)
-			progress.N = i
-			progressWriter.Set(progress)
-			if err := _ctx.Err(); err != nil {
-				progressWriter.Done()
-				return err
-			}
-			o.curKey = key
-			itemBody, ok := o.rebuilt.Tree(itemCtx, key.TreeID).ReadItem(itemCtx, key.Key)
-			if !ok {
-				o.ioErr(itemCtx, fmt.Errorf("could not read previously read item: %v", key))
+		type keyAndBody struct {
+			keyAndTree
+			Body btrfsitem.Item
+		}
+		itemChan := make(chan keyAndBody, textui.Tunable(300)) // average items-per-node≈100; let's have a buffer of ~3 nodes
+		grp := dgroup.NewGroup(stepCtx, dgroup.GroupConfig{})
+		grp.Go("io", func(stepCtx context.Context) error {
+			defer close(itemChan)
+			for _, key := range itemQueue {
+				if err := stepCtx.Err(); err != nil {
+					return err
+				}
+				itemCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.process.item", key)
+				itemBody, ok := o.rebuilt.Tree(itemCtx, key.TreeID).ReadItem(itemCtx, key.Key)
+				if !ok {
+					o.ioErr(itemCtx, fmt.Errorf("could not read previously read item: %v", key))
+				}
+				itemChan <- keyAndBody{
+					keyAndTree: key,
+					Body:       itemBody,
+				}
 			}
-			handleItem(o, itemCtx, key.TreeID, btrfstree.Item{
-				Key:  key.Key,
-				Body: itemBody,
-			})
-			if key.ItemType == btrfsitem.ROOT_ITEM_KEY {
-				o.treeQueue.Insert(key.ObjectID)
+			return nil
+		})
+		grp.Go("cpu", func(stepCtx context.Context) error {
+			defer progressWriter.Done()
+			for item := range itemChan {
+				itemCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.process.item", item.keyAndTree)
+				o.curKey = item.keyAndTree
+				handleItem(o, itemCtx, item.TreeID, btrfstree.Item{
+					Key:  item.Key,
+					Body: item.Body,
+				})
+				if item.ItemType == btrfsitem.ROOT_ITEM_KEY {
+					o.treeQueue.Insert(item.ObjectID)
+				}
+				progress.N++
+				progressWriter.Set(progress)
 			}
+			return nil
+		})
+		if err := grp.Wait(); err != nil {
+			return err
 		}
-		progress.N = len(itemQueue)
-		progressWriter.Set(progress)
-		progressWriter.Done()
 
 		// Apply augments (drain o.augmentQueue, fill o.itemQueue)
 		stepCtx = dlog.WithField(passCtx, "btrfsinspect.rebuild-nodes.rebuild.substep", "apply-augments")
-- 
cgit v1.2.3-54-g00ecf


From a433371680b2e01a5fdf05b342c9c4d9f8c3cc20 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sat, 31 Dec 2022 13:15:07 -0700
Subject: rebuildnodes/btrees: Allow leaf-node indexes to be evicted

---
 .../btrfsinspect/rebuildnodes/btrees/forrest.go    |  3 +-
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 69 ++++++++++++----------
 .../btrfsinspect/rebuildnodes/rebuild.go           |  8 +--
 lib/textui/log.go                                  | 12 ++--
 4 files changed, 49 insertions(+), 43 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 60f922d..de21d9a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -63,6 +63,7 @@ type RebuiltForrest struct {
 
 	// mutable
 	trees    containers.SyncMap[btrfsprim.ObjID, *RebuiltTree]
+	leafs    *containers.LRUCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]
 	allItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 	incItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 }
@@ -84,6 +85,7 @@ func NewRebuiltForrest(
 		cbLookupRoot: cbLookupRoot,
 		cbLookupUUID: cbLookupUUID,
 
+		leafs:    containers.NewLRUCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]](textui.Tunable(8)),
 		allItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 		incItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 	}
@@ -154,7 +156,6 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			tree.Parent, _ = ts.trees.Load(parentID)
 		}
 	}
-	tree.indexLeafs(ctx)
 
 	ts.trees.Store(treeID, tree)
 	if root != 0 {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index 1daeefc..678d25f 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -1,4 +1,4 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
@@ -31,42 +31,44 @@ type RebuiltTree struct {
 	ParentGen btrfsprim.Generation // offset of this tree's root item
 	forrest   *RebuiltForrest
 
-	// all leafs (lvl=0) that pass .isOwnerOK, whether or not they're  in the tree
-	leafToRoots map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]
-
 	// mutable
 	mu    sync.Mutex
 	Roots containers.Set[btrfsvol.LogicalAddr]
 	Leafs containers.Set[btrfsvol.LogicalAddr]
 }
 
-// initializaton (called by `RebuiltForrest.Tree()`) ///////////////////////////////////////////////////////////////////
+// .leafToRoots() //////////////////////////////////////////////////////////////////////////////////////////////////////
 
-func (tree *RebuiltTree) indexLeafs(ctx context.Context) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.substep", "index-nodes")
+// leafToRoots returns all leafs (lvl=0) in the filesystem that pass
+// .isOwnerOK, whether or not they're in the tree.
+func (tree *RebuiltTree) leafToRoots(ctx context.Context) map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr] {
+	return tree.forrest.leafs.GetOrElse(tree.ID, func() map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr] {
+		ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-nodes", fmt.Sprintf("tree=%v", tree.ID))
 
-	nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+		nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
 
-	var stats textui.Portion[int]
-	stats.D = len(tree.forrest.graph.Nodes)
-	progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	progress := func() {
-		stats.N = len(nodeToRoots)
-		progressWriter.Set(stats)
-	}
+		var stats textui.Portion[int]
+		stats.D = len(tree.forrest.graph.Nodes)
+		progressWriter := textui.NewProgress[textui.Portion[int]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+		progress := func() {
+			stats.N = len(nodeToRoots)
+			progressWriter.Set(stats)
+		}
 
-	progress()
-	for _, node := range maps.SortedKeys(tree.forrest.graph.Nodes) {
-		tree.indexNode(ctx, node, nodeToRoots, progress, nil)
-	}
-	progressWriter.Done()
+		progress()
+		for _, node := range maps.SortedKeys(tree.forrest.graph.Nodes) {
+			tree.indexNode(ctx, node, nodeToRoots, progress, nil)
+		}
+		progressWriter.Done()
 
-	tree.leafToRoots = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
-	for node, roots := range nodeToRoots {
-		if tree.forrest.graph.Nodes[node].Level == 0 && len(roots) > 0 {
-			tree.leafToRoots[node] = roots
+		ret := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+		for node, roots := range nodeToRoots {
+			if tree.forrest.graph.Nodes[node].Level == 0 && len(roots) > 0 {
+				ret[node] = roots
+			}
 		}
-	}
+		return ret
+	})
 }
 
 func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
@@ -145,13 +147,14 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	tree.Roots.Insert(rootNode)
 
 	var stats rootStats
-	stats.Leafs.D = len(tree.leafToRoots)
+	leafToRoots := tree.leafToRoots(ctx)
+	stats.Leafs.D = len(leafToRoots)
 	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	for i, leaf := range maps.SortedKeys(tree.leafToRoots) {
+	for i, leaf := range maps.SortedKeys(leafToRoots) {
 		stats.Leafs.N = i
 		progressWriter.Set(stats)
 
-		if tree.Leafs.Has(leaf) || !tree.leafToRoots[leaf].Has(rootNode) {
+		if tree.Leafs.Has(leaf) || !leafToRoots[leaf].Has(rootNode) {
 			continue
 		}
 
@@ -165,7 +168,7 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 			progressWriter.Set(stats)
 		}
 	}
-	stats.Leafs.N = len(tree.leafToRoots)
+	stats.Leafs.N = len(leafToRoots)
 	progressWriter.Set(stats)
 	progressWriter.Done()
 }
@@ -188,7 +191,7 @@ func (tree *RebuiltTree) Items(ctx context.Context) *containers.SortedMap[btrfsp
 // RebuiltTree's internal map!
 func (tree *RebuiltTree) PotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-all-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, tree.forrest.allItems, maps.SortedKeys(tree.leafToRoots))
+	return tree.items(ctx, tree.forrest.allItems, maps.SortedKeys(tree.leafToRoots(ctx)))
 }
 
 type itemIndex struct {
@@ -317,13 +320,15 @@ func (tree *RebuiltTree) ReadItem(ctx context.Context, key btrfsprim.Key) (item
 
 // LeafToRoots returns the list of potential roots (to pass to
 // .AddRoot) that include a given leaf-node.
-func (tree *RebuiltTree) LeafToRoots(leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
+func (tree *RebuiltTree) LeafToRoots(ctx context.Context, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
 	if tree.forrest.graph.Nodes[leaf].Level != 0 {
 		panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): not a leaf",
 			tree.ID, leaf))
 	}
+	tree.mu.Lock()
+	defer tree.mu.Unlock()
 	ret := make(containers.Set[btrfsvol.LogicalAddr])
-	for root := range tree.leafToRoots[leaf] {
+	for root := range tree.leafToRoots(ctx)[leaf] {
 		if tree.Roots.Has(root) {
 			panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): tree contains root=%v but not leaf",
 				tree.ID, leaf, root))
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index cefa668..3696a8a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -447,7 +447,7 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, objID btr
 	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).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.Tree(ctx, treeID).LeafToRoots(v.Node))
+			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -484,7 +484,7 @@ func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, tgt bt
 	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { return tgt.Cmp(k) },
 		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(v.Node))
+			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, v.Node))
 			return true
 		})
 	o.wantAugment(ctx, treeID, wants)
@@ -526,7 +526,7 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, objID
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Cmp(k) },
 		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
 			if fn(v) {
-				wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(v.Node))
+				wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, v.Node))
 			}
 			return true
 		})
@@ -706,7 +706,7 @@ func (o *rebuilder) _wantRange(
 				}
 				wantCtx := dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key",
 					fmt.Sprintf("tree=%v key={%v %v %v-%v}", treeID, objID, typ, gap.Beg, gap.End))
-				o.wantAugment(wantCtx, treeID, o.rebuilt.Tree(wantCtx, treeID).LeafToRoots(v.Node))
+				o.wantAugment(wantCtx, treeID, o.rebuilt.Tree(wantCtx, treeID).LeafToRoots(wantCtx, v.Node))
 				last = runEnd
 
 				return true
diff --git a/lib/textui/log.go b/lib/textui/log.go
index da303dd..8327b67 100644
--- a/lib/textui/log.go
+++ b/lib/textui/log.go
@@ -322,17 +322,17 @@ func fieldOrd(key string) int {
 		return -8
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree":
 		return -7
-	case "btrfsinspect.rebuild-nodes.rebuild.add-tree.substep":
-		return -6
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.key":
-		return -5
+		return -6
 	case "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.reason":
-		return -4
+		return -5
 	case "btrfsinspect.rebuild-nodes.rebuild.add-root":
-		return -3
+		return -4
 	case "btrfsinspect.rebuild-nodes.rebuild.index-inc-items":
-		return -2
+		return -3
 	case "btrfsinspect.rebuild-nodes.rebuild.index-all-items":
+		return -2
+	case "btrfsinspect.rebuild-nodes.rebuild.index-nodes":
 		return -1
 
 	// other ///////////////////////////////////////////////////////////////
-- 
cgit v1.2.3-54-g00ecf


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/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go')

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-54-g00ecf


From 6f7a95e0952887c02421431bdd2a879387deebe0 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 3 Jan 2023 19:15:34 -0700
Subject: rebuildnodes/btrees: Touch up a comment

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index c9747ff..7c64bf0 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -290,9 +290,9 @@ func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bo
 			// Retain the old higher-gen one.
 			return false
 		default:
-			// This is a panic because I'm not really sure what the best way to
-			// handle this is, and so if this happens I want the program to crash
-			// and force me to figure out how to handle it.
+			// TODO: This is a panic because I'm not really sure what the
+			// best way to handle this is, and so if this happens I want the
+			// program to crash and force me to figure out how to handle it.
 			panic(fmt.Errorf("dup nodes in tree=%v: old=%v=%v ; new=%v=%v",
 				tree.ID,
 				oldNode, tree.forrest.graph.Nodes[oldNode],
-- 
cgit v1.2.3-54-g00ecf


From a794062b5391b83e9f932c06f3953964e8e70b68 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 3 Jan 2023 19:22:20 -0700
Subject: rebuildnodes/btrees: Move code up/down in the file, add comments

---
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 126 +++++++++++----------
 1 file changed, 66 insertions(+), 60 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index 7c64bf0..ae4065d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -35,9 +35,15 @@ type RebuiltTree struct {
 	mu    sync.Mutex
 	Roots containers.Set[btrfsvol.LogicalAddr]
 	Leafs containers.Set[btrfsvol.LogicalAddr]
+	// There are 3 more mutable "members" that are protected by
+	// `mu`; but they live in a shared LRUcache:
+	//
+	//  1. tree.leafToRoots()    = tree.forrest.leafs.Load(tree.ID)
+	//  2. tree.PotentialItems() = tree.forrest.allItems.Load(tree.ID)
+	//  3. tree.Items()          = tree.forrest.incItems.Load(tree.ID)
 }
 
-// .leafToRoots() //////////////////////////////////////////////////////////////////////////////////////////////////////
+// LRU member 1: .leafToRoots() ////////////////////////////////////////////////////////////////////////////////////////
 
 // leafToRoots returns all leafs (lvl=0) in the filesystem that pass
 // .isOwnerOK, whether or not they're in the tree.
@@ -124,65 +130,7 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 	}
 }
 
-// .AddRoot() //////////////////////////////////////////////////////////////////////////////////////////////////////////
-
-type rootStats struct {
-	Leafs      textui.Portion[int]
-	AddedLeafs int
-	AddedItems int
-}
-
-func (s rootStats) String() string {
-	return textui.Sprintf("%v (added %v leafs, added %v items)",
-		s.Leafs, s.AddedLeafs, s.AddedItems)
-}
-
-// AddRoot adds an additional root node to the tree.  It is useful to
-// call .AddRoot() to re-attach part of the tree that has been broken
-// off.
-func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
-	tree.mu.Lock()
-	defer tree.mu.Unlock()
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
-	tree.Roots.Insert(rootNode)
-
-	var stats rootStats
-	leafToRoots := tree.leafToRoots(ctx)
-	stats.Leafs.D = len(leafToRoots)
-	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	for i, leaf := range maps.SortedKeys(leafToRoots) {
-		stats.Leafs.N = i
-		progressWriter.Set(stats)
-
-		if tree.Leafs.Has(leaf) || !leafToRoots[leaf].Has(rootNode) {
-			continue
-		}
-
-		tree.Leafs.Insert(leaf)
-		stats.AddedLeafs++
-		progressWriter.Set(stats)
-
-		for _, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
-			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
-			stats.AddedItems++
-			progressWriter.Set(stats)
-		}
-	}
-	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() //////////////////////////////////////////////////////////////////////////////////////
+// LRU members 2 and 3: .Items() and .PotentialItems() /////////////////////////////////////////////////////////////////
 
 // Items returns a map of the items contained in this tree.
 //
@@ -301,6 +249,64 @@ func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bo
 	}
 }
 
+// .AddRoot() //////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+type rootStats struct {
+	Leafs      textui.Portion[int]
+	AddedLeafs int
+	AddedItems int
+}
+
+func (s rootStats) String() string {
+	return textui.Sprintf("%v (added %v leafs, added %v items)",
+		s.Leafs, s.AddedLeafs, s.AddedItems)
+}
+
+// AddRoot adds an additional root node to the tree.  It is useful to
+// call .AddRoot() to re-attach part of the tree that has been broken
+// off.
+func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
+	tree.mu.Lock()
+	defer tree.mu.Unlock()
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
+	tree.Roots.Insert(rootNode)
+
+	var stats rootStats
+	leafToRoots := tree.leafToRoots(ctx)
+	stats.Leafs.D = len(leafToRoots)
+	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	for i, leaf := range maps.SortedKeys(leafToRoots) {
+		stats.Leafs.N = i
+		progressWriter.Set(stats)
+
+		if tree.Leafs.Has(leaf) || !leafToRoots[leaf].Has(rootNode) {
+			continue
+		}
+
+		tree.Leafs.Insert(leaf)
+		stats.AddedLeafs++
+		progressWriter.Set(stats)
+
+		for _, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
+			stats.AddedItems++
+			progressWriter.Set(stats)
+		}
+	}
+	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
+		})
+	}
+}
+
 // main public API /////////////////////////////////////////////////////////////////////////////////////////////////////
 
 // COWDistance returns how many COW-snapshots down the 'tree' is from
-- 
cgit v1.2.3-54-g00ecf


From 94e662ff74a23bf579eb2bcf6f2b152b13436a01 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 3 Jan 2023 20:56:20 -0700
Subject: rebuildnodes/btrees: Rework to avoid the .Leafs member

Save some memory.
---
 .../btrfsinspect/rebuildnodes/btrees/forrest.go    |  1 -
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 34 ++++++++++++++++------
 2 files changed, 25 insertions(+), 10 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 3ed27b6..26fa64e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -126,7 +126,6 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 	tree := &RebuiltTree{
 		ID:      treeID,
 		Roots:   make(containers.Set[btrfsvol.LogicalAddr]),
-		Leafs:   make(containers.Set[btrfsvol.LogicalAddr]),
 		forrest: ts,
 	}
 	var root btrfsvol.LogicalAddr
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index ae4065d..a25825a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -34,9 +34,10 @@ type RebuiltTree struct {
 	// mutable
 	mu    sync.Mutex
 	Roots containers.Set[btrfsvol.LogicalAddr]
-	Leafs containers.Set[btrfsvol.LogicalAddr]
 	// There are 3 more mutable "members" that are protected by
-	// `mu`; but they live in a shared LRUcache:
+	// `mu`; but they live in a shared LRUcache.  They are all
+	// derived from tree.Roots, which is why it's OK if they get
+	// evicted.
 	//
 	//  1. tree.leafToRoots()    = tree.forrest.leafs.Load(tree.ID)
 	//  2. tree.PotentialItems() = tree.forrest.allItems.Load(tree.ID)
@@ -138,7 +139,7 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 // RebuiltTree's internal map!
 func (tree *RebuiltTree) Items(ctx context.Context) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-inc-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, tree.forrest.incItems, maps.SortedKeys(tree.Leafs))
+	return tree.items(ctx, tree.forrest.incItems, tree.Roots.HasAny)
 }
 
 // PotentialItems returns a map of items that could be added to this
@@ -148,7 +149,10 @@ func (tree *RebuiltTree) Items(ctx context.Context) *containers.SortedMap[btrfsp
 // RebuiltTree's internal map!
 func (tree *RebuiltTree) PotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-all-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, tree.forrest.allItems, maps.SortedKeys(tree.leafToRoots(ctx)))
+	return tree.items(ctx, tree.forrest.allItems,
+		func(_ containers.Set[btrfsvol.LogicalAddr]) bool {
+			return true
+		})
 }
 
 type itemIndex struct {
@@ -168,9 +172,20 @@ func (s itemStats) String() string {
 		s.Leafs, s.NumItems, s.NumDups)
 }
 
-func (tree *RebuiltTree) items(ctx context.Context, cache *containers.LRUCache[btrfsprim.ObjID, *itemIndex], leafs []btrfsvol.LogicalAddr) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
+func (tree *RebuiltTree) items(ctx context.Context, cache *containers.LRUCache[btrfsprim.ObjID, *itemIndex],
+	leafFn func(roots containers.Set[btrfsvol.LogicalAddr]) bool,
+) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
 	tree.mu.Lock()
 	defer tree.mu.Unlock()
+
+	var leafs []btrfsvol.LogicalAddr
+	for leaf, roots := range tree.leafToRoots(ctx) {
+		if leafFn(roots) {
+			leafs = append(leafs, leaf)
+		}
+	}
+	slices.Sort(leafs)
+
 	index := cache.GetOrElse(tree.ID, func() *itemIndex {
 		return &itemIndex{
 			Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
@@ -269,21 +284,20 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	tree.mu.Lock()
 	defer tree.mu.Unlock()
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
-	tree.Roots.Insert(rootNode)
 
-	var stats rootStats
 	leafToRoots := tree.leafToRoots(ctx)
+
+	var stats rootStats
 	stats.Leafs.D = len(leafToRoots)
 	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 	for i, leaf := range maps.SortedKeys(leafToRoots) {
 		stats.Leafs.N = i
 		progressWriter.Set(stats)
 
-		if tree.Leafs.Has(leaf) || !leafToRoots[leaf].Has(rootNode) {
+		if tree.Roots.HasAny(leafToRoots[leaf]) || !leafToRoots[leaf].Has(rootNode) {
 			continue
 		}
 
-		tree.Leafs.Insert(leaf)
 		stats.AddedLeafs++
 		progressWriter.Set(stats)
 
@@ -297,6 +311,8 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	progressWriter.Set(stats)
 	progressWriter.Done()
 
+	tree.Roots.Insert(rootNode)
+
 	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 {
-- 
cgit v1.2.3-54-g00ecf


From c7d387f5ddd39e2d359c2c0c2ef52536ab650160 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 3 Jan 2023 21:03:25 -0700
Subject: rebuildnodes/btrees: Don't include .Items() in .PotentialItems()

Save some memory.
---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go |  4 ++--
 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go    | 10 +++++-----
 lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go        |  8 ++++----
 3 files changed, 11 insertions(+), 11 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 26fa64e..ff6b1c5 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -64,8 +64,8 @@ type RebuiltForrest struct {
 	// mutable
 	trees    containers.SyncMap[btrfsprim.ObjID, *RebuiltTree]
 	leafs    *containers.LRUCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]
-	allItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 	incItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
+	excItems *containers.LRUCache[btrfsprim.ObjID, *itemIndex]
 }
 
 // NewRebuiltForrest returns a new RebuiltForrest instance.  All of
@@ -86,8 +86,8 @@ func NewRebuiltForrest(
 		cbLookupUUID: cbLookupUUID,
 
 		leafs:    containers.NewLRUCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]](textui.Tunable(8)),
-		allItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 		incItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
+		excItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 	}
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index a25825a..65f76b2 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -40,8 +40,8 @@ type RebuiltTree struct {
 	// evicted.
 	//
 	//  1. tree.leafToRoots()    = tree.forrest.leafs.Load(tree.ID)
-	//  2. tree.PotentialItems() = tree.forrest.allItems.Load(tree.ID)
-	//  3. tree.Items()          = tree.forrest.incItems.Load(tree.ID)
+	//  2. tree.Items()          = tree.forrest.incItems.Load(tree.ID)
+	//  3. tree.PotentialItems() = tree.forrest.excItems.Load(tree.ID)
 }
 
 // LRU member 1: .leafToRoots() ////////////////////////////////////////////////////////////////////////////////////////
@@ -149,9 +149,9 @@ func (tree *RebuiltTree) Items(ctx context.Context) *containers.SortedMap[btrfsp
 // RebuiltTree's internal map!
 func (tree *RebuiltTree) PotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.index-all-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, tree.forrest.allItems,
-		func(_ containers.Set[btrfsvol.LogicalAddr]) bool {
-			return true
+	return tree.items(ctx, tree.forrest.excItems,
+		func(roots containers.Set[btrfsvol.LogicalAddr]) bool {
+			return !tree.Roots.HasAny(roots)
 		})
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index fbbda26..4709db2 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -723,8 +723,8 @@ func (o *rebuilder) _wantRange(
 		return
 	}
 
-	sizeFn := func(key btrfsprim.Key) (uint64, error) {
-		ptr, ok := o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Load(key)
+	sizeFn := func(items *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr], key btrfsprim.Key) (uint64, error) {
+		ptr, ok := items.Load(key)
 		if !ok {
 			panic(fmt.Errorf("should not happen: could not load key: %v", keyAndTree{TreeID: treeID, Key: key}))
 		}
@@ -776,7 +776,7 @@ func (o *rebuilder) _wantRange(
 			}
 		},
 		func(runKey btrfsprim.Key, _ keyio.ItemPtr) bool {
-			runSize, err := sizeFn(runKey)
+			runSize, err := sizeFn(o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx), runKey)
 			if err != nil {
 				o.fsErr(ctx, fmt.Errorf("get size: %v: %w", keyAndTree{TreeID: treeID, Key: runKey}, err))
 				return true
@@ -848,7 +848,7 @@ func (o *rebuilder) _wantRange(
 				}
 			},
 			func(k btrfsprim.Key, v keyio.ItemPtr) bool {
-				runSize, err := sizeFn(k)
+				runSize, err := sizeFn(o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx), k)
 				if err != nil {
 					o.fsErr(ctx, fmt.Errorf("get size: %v: %w", keyAndTree{TreeID: treeID, Key: k}, err))
 					return true
-- 
cgit v1.2.3-54-g00ecf


From 98de12d606765ef21a1aae10abfb4c33a1a9f23b Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 5 Jan 2023 00:30:41 -0700
Subject: rebuildnodes/btrees: Rework RebuiltTree.items() to be faster in the
 happy-path

And also have the cache consume less memory.
---
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 88 ++++++++++------------
 1 file changed, 38 insertions(+), 50 deletions(-)

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index 65f76b2..cd05911 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -155,11 +155,7 @@ func (tree *RebuiltTree) PotentialItems(ctx context.Context) *containers.SortedM
 		})
 }
 
-type itemIndex struct {
-	NumDups int
-	Leafs   containers.Set[btrfsvol.LogicalAddr]
-	Items   containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
-}
+type itemIndex = containers.SortedMap[btrfsprim.Key, keyio.ItemPtr]
 
 type itemStats struct {
 	Leafs    textui.Portion[int]
@@ -178,58 +174,48 @@ func (tree *RebuiltTree) items(ctx context.Context, cache *containers.LRUCache[b
 	tree.mu.Lock()
 	defer tree.mu.Unlock()
 
-	var leafs []btrfsvol.LogicalAddr
-	for leaf, roots := range tree.leafToRoots(ctx) {
-		if leafFn(roots) {
-			leafs = append(leafs, leaf)
-		}
-	}
-	slices.Sort(leafs)
-
-	index := cache.GetOrElse(tree.ID, func() *itemIndex {
-		return &itemIndex{
-			Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
+	return cache.GetOrElse(tree.ID, func() *itemIndex {
+		var leafs []btrfsvol.LogicalAddr
+		for leaf, roots := range tree.leafToRoots(ctx) {
+			if leafFn(roots) {
+				leafs = append(leafs, leaf)
+			}
 		}
-	})
+		slices.Sort(leafs)
 
-	var stats itemStats
-	stats.Leafs.D = len(leafs)
-	progressWriter := textui.NewProgress[itemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	progress := func(doneLeafs int) {
-		stats.Leafs.N = doneLeafs
-		stats.NumItems = index.Items.Len()
-		stats.NumDups = index.NumDups
-		progressWriter.Set(stats)
-	}
+		var stats itemStats
+		stats.Leafs.D = len(leafs)
+		progressWriter := textui.NewProgress[itemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 
-	for i, leaf := range leafs {
-		if index.Leafs.Has(leaf) {
-			continue
-		}
-		progress(i)
-		index.Leafs.Insert(leaf)
-		for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
-			newPtr := keyio.ItemPtr{
-				Node: leaf,
-				Idx:  j,
-			}
-			if oldPtr, exists := index.Items.Load(itemKey); !exists {
-				index.Items.Store(itemKey, newPtr)
-			} else {
-				index.NumDups++
-				if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
-					index.Items.Store(itemKey, newPtr)
+		index := new(containers.SortedMap[btrfsprim.Key, keyio.ItemPtr])
+		for i, leaf := range leafs {
+			stats.Leafs.N = i
+			progressWriter.Set(stats)
+			for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+				newPtr := keyio.ItemPtr{
+					Node: leaf,
+					Idx:  j,
+				}
+				if oldPtr, exists := index.Load(itemKey); !exists {
+					index.Store(itemKey, newPtr)
+					stats.NumItems++
+				} else {
+					if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
+						index.Store(itemKey, newPtr)
+					}
+					stats.NumDups++
 				}
+				progressWriter.Set(stats)
 			}
-			progress(i)
 		}
-	}
-	if stats.Leafs.N > 0 {
-		progress(len(leafs))
-		progressWriter.Done()
-	}
+		if stats.Leafs.N > 0 {
+			stats.Leafs.N = len(leafs)
+			progressWriter.Set(stats)
+			progressWriter.Done()
+		}
 
-	return &index.Items
+		return index
+	})
 }
 
 func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bool {
@@ -312,6 +298,8 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	progressWriter.Done()
 
 	tree.Roots.Insert(rootNode)
+	tree.forrest.incItems.Remove(tree.ID) // force re-gen
+	tree.forrest.excItems.Remove(tree.ID) // force re-gen
 
 	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 {
-- 
cgit v1.2.3-54-g00ecf


From 3b385f26973e45b4c2e2f3ebf9d52ab0131cff5e Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Thu, 5 Jan 2023 00:47:18 -0700
Subject: rebuildnodes/btrees: Switch from a sync.Mutex to a sync.RWMutex

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

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

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index cd05911..c381274 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -32,7 +32,7 @@ type RebuiltTree struct {
 	forrest   *RebuiltForrest
 
 	// mutable
-	mu    sync.Mutex
+	mu    sync.RWMutex
 	Roots containers.Set[btrfsvol.LogicalAddr]
 	// There are 3 more mutable "members" that are protected by
 	// `mu`; but they live in a shared LRUcache.  They are all
@@ -171,8 +171,8 @@ func (s itemStats) String() string {
 func (tree *RebuiltTree) items(ctx context.Context, cache *containers.LRUCache[btrfsprim.ObjID, *itemIndex],
 	leafFn func(roots containers.Set[btrfsvol.LogicalAddr]) bool,
 ) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
-	tree.mu.Lock()
-	defer tree.mu.Unlock()
+	tree.mu.RLock()
+	defer tree.mu.RUnlock()
 
 	return cache.GetOrElse(tree.ID, func() *itemIndex {
 		var leafs []btrfsvol.LogicalAddr
@@ -344,8 +344,8 @@ func (tree *RebuiltTree) LeafToRoots(ctx context.Context, leaf btrfsvol.LogicalA
 		panic(fmt.Errorf("should not happen: (tree=%v).LeafToRoots(leaf=%v): not a leaf",
 			tree.ID, leaf))
 	}
-	tree.mu.Lock()
-	defer tree.mu.Unlock()
+	tree.mu.RLock()
+	defer tree.mu.RUnlock()
 	ret := make(containers.Set[btrfsvol.LogicalAddr])
 	for root := range tree.leafToRoots(ctx)[leaf] {
 		if tree.Roots.Has(root) {
-- 
cgit v1.2.3-54-g00ecf