From 5590315b11557aa48999d09700631ad9239ce03b Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Tue, 27 Dec 2022 01:14:34 -0700
Subject: rebuildnodes: Optimize: Try to avoid disk access for DIR_INDEX

---
 .../btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go   | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index b53a28e..b1ae7be 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -363,15 +363,23 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd
 	}
 }
 
+// Resolve a key to a keyio.ItemPtr.
+//
+// It is not nescessary to call AddTree for that tree first; Resolve will
+// call it for you.
+func (ts *RebuiltTrees) Resolve(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
+	if !ts.AddTree(ctx, treeID) {
+		return keyio.ItemPtr{}, false
+	}
+	return ts.trees[treeID].Items.Load(key)
+}
+
 // Load reads an item from a tree.
 //
 // It is not nescessary to call AddTree for that tree first; Load will
 // call it for you.
 func (ts *RebuiltTrees) Load(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
-	if !ts.AddTree(ctx, treeID) {
-		return nil, false
-	}
-	ptr, ok := ts.trees[treeID].Items.Load(key)
+	ptr, ok := ts.Resolve(ctx, treeID, key)
 	if !ok {
 		return nil, false
 	}
-- 
cgit v1.2.3-2-g168b


From 8efc82d0b1a167830970135c78d173667080b116 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 28 Dec 2022 18:49:09 -0700
Subject: rebuildnodes: Support graceful shutdown

---
 .../btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go           | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index b1ae7be..ffe225a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -305,7 +305,7 @@ func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
 
 	progress()
 	for _, node := range maps.SortedKeys(graph.Nodes) {
-		tree.indexNode(graph, node, nodeToRoots, progress, nil)
+		tree.indexNode(ctx, graph, node, nodeToRoots, progress, nil)
 	}
 	progressWriter.Done()
 
@@ -317,8 +317,11 @@ func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
 	}
 }
 
-func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
+func (tree *rebuiltTree) indexNode(ctx context.Context, graph pkggraph.Graph, 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
 	}
@@ -339,7 +342,7 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd
 		return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner, graph.Nodes[kp.FromNode].Generation)
 	})
 	for _, kp := range kps {
-		tree.indexNode(graph, kp.FromNode, index, progress, stack)
+		tree.indexNode(ctx, graph, kp.FromNode, index, progress, stack)
 		if len(index[kp.FromNode]) > 0 {
 			if roots == nil {
 				roots = make(containers.Set[btrfsvol.LogicalAddr])
-- 
cgit v1.2.3-2-g168b


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/forrest.go    | 173 ++++++++
 .../rebuildnodes/btrees/rebuilt_btrees.go          | 482 ---------------------
 .../btrfsinspect/rebuildnodes/btrees/tree.go       | 299 +++++++++++++
 3 files changed, 472 insertions(+), 482 deletions(-)
 create mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
 delete mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
 create mode 100644 lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
new file mode 100644
index 0000000..791e646
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -0,0 +1,173 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrees
+
+import (
+	"context"
+
+	"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/btrfstree"
+	"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/slices"
+)
+
+// RebuiltForrest is an abstraction for rebuilding and accessing
+// potentially broken btrees.
+//
+// It is conceptually a btrfstree.TreeOperator, and adds similar
+// broken-tree handling to btrfsutil.BrokenForrest.  However, the API
+// is different thant btrfstree.TreeOperator, and is much more
+// efficient than btrfsutil.BrokenForrest.
+//
+// The efficiency improvements are possible because of the API
+// differences, which are necessary for how it is used in
+// rebuildnodes:
+//
+//   - it consumes an already-read graph.Graph instead of reading the
+//     graph itself
+//
+//   - it does not use `btrfstree.TreePath`
+//
+//   - it does not keep track of errors encountered in a tree
+//
+// Additionally, it provides some functionality that
+// btrfsutil.BrokenForrest does not:
+//
+//   - it provides a .LeafToRoots() method to advise on what
+//     additional roots should be added
+//
+//   - it provides a .COWDistance() method to compare how related two
+//     trees are
+//
+// A zero RebuiltForrest is invalid; it must be initialized with
+// NewRebuiltForrest().
+type RebuiltForrest struct {
+	// static
+	sb    btrfstree.Superblock
+	graph pkggraph.Graph
+	keyIO *keyio.Handle
+
+	// static callbacks
+	cbAddedItem  func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
+	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool)
+	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
+
+	// mutable
+	trees map[btrfsprim.ObjID]*RebuiltTree
+}
+
+// NewRebuiltForrest returns a new RebuiltForrest instance.  All of
+// the callbacks must be non-nil.
+func NewRebuiltForrest(
+	sb btrfstree.Superblock, graph pkggraph.Graph, keyIO *keyio.Handle,
+	cbAddedItem func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key),
+	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool),
+	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
+) *RebuiltForrest {
+	return &RebuiltForrest{
+		sb:    sb,
+		graph: graph,
+		keyIO: keyIO,
+
+		cbAddedItem:  cbAddedItem,
+		cbLookupRoot: cbLookupRoot,
+		cbLookupUUID: cbLookupUUID,
+
+		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
+	}
+}
+
+// Tree returns a given tree, initializing it if nescessary.  If it is
+// unable to initialize the tree, then nil is returned, and nothing is
+// done to the forrest.
+//
+// The tree is initialized with the normal root node of the tree.
+func (ts *RebuiltForrest) Tree(ctx context.Context, treeID btrfsprim.ObjID) *RebuiltTree {
+	if !ts.addTree(ctx, treeID, nil) {
+		return nil
+	}
+	return ts.trees[treeID]
+}
+
+func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, stack []btrfsprim.ObjID) (ok bool) {
+	if _, ok := ts.trees[treeID]; ok {
+		return true
+	}
+	if slices.Contains(treeID, stack) {
+		return false
+	}
+	stack = append(stack, treeID)
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree", stack)
+	dlog.Info(ctx, "adding tree...")
+
+	tree := &RebuiltTree{
+		ID:      treeID,
+		Roots:   make(containers.Set[btrfsvol.LogicalAddr]),
+		Leafs:   make(containers.Set[btrfsvol.LogicalAddr]),
+		forrest: ts,
+	}
+	var root btrfsvol.LogicalAddr
+	switch treeID {
+	case btrfsprim.ROOT_TREE_OBJECTID:
+		root = ts.sb.RootTree
+	case btrfsprim.CHUNK_TREE_OBJECTID:
+		root = ts.sb.ChunkTree
+	case btrfsprim.TREE_LOG_OBJECTID:
+		root = ts.sb.LogTree
+	case btrfsprim.BLOCK_GROUP_TREE_OBJECTID:
+		root = ts.sb.BlockGroupRoot
+	default:
+		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
+			return false
+		}
+		rootOff, rootItem, ok := ts.cbLookupRoot(ctx, treeID)
+		if !ok {
+			return false
+		}
+		root = rootItem.ByteNr
+		tree.UUID = rootItem.UUID
+		if rootItem.ParentUUID != (btrfsprim.UUID{}) {
+			tree.ParentGen = rootOff
+			if !ts.addTree(ctx, btrfsprim.UUID_TREE_OBJECTID, stack) {
+				return false
+			}
+			parentID, ok := ts.cbLookupUUID(ctx, rootItem.ParentUUID)
+			if !ok {
+				return false
+			}
+			if !ts.addTree(ctx, parentID, append(stack, treeID)) {
+				return false
+			}
+			tree.Parent = ts.trees[parentID]
+		}
+	}
+	tree.indexLeafs(ctx)
+
+	ts.trees[treeID] = tree
+	if root != 0 {
+		tree.AddRoot(ctx, root)
+	}
+
+	return true
+}
+
+// ListRoots returns a listing of all initialized trees and their root
+// nodes.
+//
+// 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
+	}
+	return ret
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
deleted file mode 100644
index ffe225a..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ /dev/null
@@ -1,482 +0,0 @@
-// Copyright (C) 2022-2023  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/btrfstree"
-	"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
-
-	// 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]
-}
-
-// 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
-	}
-}
-
-// 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++
-	}
-}
-
-func (tree *rebuiltTree) shouldReplace(graph pkggraph.Graph, oldNode, newNode btrfsvol.LogicalAddr) bool {
-	oldDist, _ := tree.cowDistance(graph.Nodes[oldNode].Owner)
-	newDist, _ := tree.cowDistance(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 := graph.Nodes[oldNode].Generation
-		newGen := 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, graph.Nodes[oldNode],
-				newNode, graph.Nodes[newNode]))
-		}
-	}
-}
-
-// RebuiltTrees is an abstraction for rebuilding and accessing
-// potentially broken btrees.
-//
-// It is conceptually a btrfstree.TreeOperator, and adds similar
-// broken-tree handling to btrfsutil.BrokenTrees.  However, the API is
-// different thant btrfstree.TreeOperator, and is much more efficient
-// than btrfsutil.BrokenTrees.
-//
-// The efficiency improvements are possible because of the API
-// differences, which are necessary for how it is used in
-// rebuildnodes:
-//
-//   - it consumes an already-read graph.Graph instead of reading the
-//     graph itself
-//
-//   - it does not use `btrfstree.TreePath`
-//
-//   - it does not keep track of errors encountered in a tree
-//
-// Additionally, it provides some functionality that
-// btrfsutil.BrokenTrees does not:
-//
-//   - it provides a .LeafToRoots() method to advise on what
-//     additional roots should be added
-//
-//   - it provides a .COWDistance() method to compare how related two
-//     trees are
-//
-// A zero RebuiltTrees is invalid; it must be initialized with
-// NewRebuiltTrees().
-type RebuiltTrees struct {
-	// static
-	sb    btrfstree.Superblock
-	graph pkggraph.Graph
-	keyIO *keyio.Handle
-
-	// static callbacks
-	cbAddedItem  func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
-	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool)
-	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
-
-	// mutable
-	trees map[btrfsprim.ObjID]*rebuiltTree
-}
-
-// NewRebuiltTrees returns a new RebuiltTrees instance.  All of the
-// callbacks must be non-nil.
-func NewRebuiltTrees(
-	sb btrfstree.Superblock, graph pkggraph.Graph, keyIO *keyio.Handle,
-	cbAddedItem func(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key),
-	cbLookupRoot func(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool),
-	cbLookupUUID func(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool),
-) *RebuiltTrees {
-	return &RebuiltTrees{
-		sb:    sb,
-		graph: graph,
-		keyIO: keyIO,
-
-		cbAddedItem:  cbAddedItem,
-		cbLookupRoot: cbLookupRoot,
-		cbLookupUUID: cbLookupUUID,
-
-		trees: make(map[btrfsprim.ObjID]*rebuiltTree),
-	}
-}
-
-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 an existing tree.  It is
-// useful to call .AddRoot() to re-attach part of the tree that has
-// been broken off.
-//
-// It is invalid (panic) to call AddRoot for a tree without having
-// called AddTree first.
-func (ts *RebuiltTrees) AddRoot(ctx context.Context, treeID btrfsprim.ObjID, rootNode btrfsvol.LogicalAddr) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-root", fmt.Sprintf("tree=%v rootNode=%v", treeID, rootNode))
-	tree := ts.trees[treeID]
-	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 ts.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(ts.graph, oldPtr.Node, newPtr.Node) {
-				tree.Items.Store(itemKey, newPtr)
-				stats.ReplacedItems++
-			}
-			ts.cbAddedItem(ctx, treeID, itemKey)
-			progressWriter.Set(stats)
-		}
-	}
-	stats.Leafs.N = len(tree.leafToRoots)
-	progressWriter.Set(stats)
-	progressWriter.Done()
-}
-
-// AddTree initializes the given tree, returning true if it was able
-// to do so, or false if there was a problem and nothing was done.
-// The tree is initialized with the normal root node of the tree.
-//
-// Subsequent calls to AddTree for the same tree are no-ops.
-func (ts *RebuiltTrees) AddTree(ctx context.Context, treeID btrfsprim.ObjID) (ok bool) {
-	return ts.addTree(ctx, treeID, nil)
-}
-
-func (ts *RebuiltTrees) addTree(ctx context.Context, treeID btrfsprim.ObjID, stack []btrfsprim.ObjID) (ok bool) {
-	if _, ok := ts.trees[treeID]; ok {
-		return true
-	}
-	if slices.Contains(treeID, stack) {
-		return false
-	}
-	stack = append(stack, treeID)
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree", stack)
-	dlog.Info(ctx, "adding tree...")
-
-	tree := &rebuiltTree{
-		ID:    treeID,
-		Roots: make(containers.Set[btrfsvol.LogicalAddr]),
-		Leafs: make(containers.Set[btrfsvol.LogicalAddr]),
-	}
-	var root btrfsvol.LogicalAddr
-	switch treeID {
-	case btrfsprim.ROOT_TREE_OBJECTID:
-		root = ts.sb.RootTree
-	case btrfsprim.CHUNK_TREE_OBJECTID:
-		root = ts.sb.ChunkTree
-	case btrfsprim.TREE_LOG_OBJECTID:
-		root = ts.sb.LogTree
-	case btrfsprim.BLOCK_GROUP_TREE_OBJECTID:
-		root = ts.sb.BlockGroupRoot
-	default:
-		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
-			return false
-		}
-		rootOff, rootItem, ok := ts.cbLookupRoot(ctx, treeID)
-		if !ok {
-			return false
-		}
-		root = rootItem.ByteNr
-		tree.UUID = rootItem.UUID
-		if rootItem.ParentUUID != (btrfsprim.UUID{}) {
-			tree.ParentGen = rootOff
-			if !ts.addTree(ctx, btrfsprim.UUID_TREE_OBJECTID, stack) {
-				return false
-			}
-			parentID, ok := ts.cbLookupUUID(ctx, rootItem.ParentUUID)
-			if !ok {
-				return false
-			}
-			if !ts.addTree(ctx, parentID, append(stack, treeID)) {
-				return false
-			}
-			tree.Parent = ts.trees[parentID]
-		}
-	}
-	tree.indexLeafs(ctx, ts.graph)
-
-	ts.trees[treeID] = tree
-	if root != 0 {
-		ts.AddRoot(ctx, treeID, root)
-	}
-
-	return true
-}
-
-func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
-	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(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(graph.Nodes) {
-		tree.indexNode(ctx, graph, node, nodeToRoots, progress, nil)
-	}
-	progressWriter.Done()
-
-	tree.leafToRoots = make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
-	for node, roots := range nodeToRoots {
-		if graph.Nodes[node].Level == 0 && len(roots) > 0 {
-			tree.leafToRoots[node] = roots
-		}
-	}
-}
-
-func (tree *rebuiltTree) indexNode(ctx context.Context, graph pkggraph.Graph, 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 graph.FinalCheck() should
-		// have already checked for loops.
-		panic("loop")
-	}
-	if !tree.isOwnerOK(graph.Nodes[node].Owner, graph.Nodes[node].Generation) {
-		index[node] = nil
-		return
-	}
-
-	// tree.leafToRoots
-	stack = append(stack, node)
-	var roots containers.Set[btrfsvol.LogicalAddr]
-	kps := slices.RemoveAllFunc(graph.EdgesTo[node], func(kp *pkggraph.Edge) bool {
-		return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner, graph.Nodes[kp.FromNode].Generation)
-	})
-	for _, kp := range kps {
-		tree.indexNode(ctx, graph, 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 graph.Nodes[node].Items {
-		if oldPtr, ok := tree.keys.Load(key); !ok || tree.shouldReplace(graph, oldPtr.Node, node) {
-			tree.keys.Store(key, keyio.ItemPtr{
-				Node: node,
-				Idx:  i,
-			})
-		}
-	}
-}
-
-// Resolve a key to a keyio.ItemPtr.
-//
-// It is not nescessary to call AddTree for that tree first; Resolve will
-// call it for you.
-func (ts *RebuiltTrees) Resolve(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (ptr keyio.ItemPtr, ok bool) {
-	if !ts.AddTree(ctx, treeID) {
-		return keyio.ItemPtr{}, false
-	}
-	return ts.trees[treeID].Items.Load(key)
-}
-
-// Load reads an item from a tree.
-//
-// It is not nescessary to call AddTree for that tree first; Load will
-// call it for you.
-func (ts *RebuiltTrees) Load(ctx context.Context, treeID btrfsprim.ObjID, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
-	ptr, ok := ts.Resolve(ctx, treeID, key)
-	if !ok {
-		return nil, false
-	}
-	return ts.keyIO.ReadItem(ctx, ptr)
-}
-
-// Search searches for an item from a tree.
-//
-// It is not nescessary to call AddTree for that tree first; Search
-// will call it for you.
-func (ts *RebuiltTrees) Search(ctx context.Context, treeID btrfsprim.ObjID, fn func(btrfsprim.Key) int) (key btrfsprim.Key, ok bool) {
-	if !ts.AddTree(ctx, treeID) {
-		return btrfsprim.Key{}, false
-	}
-	k, _, ok := ts.trees[treeID].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.
-//
-// It is not nescessary to call AddTree for that tree first; SearchAll
-// will call it for you.
-func (ts *RebuiltTrees) SearchAll(ctx context.Context, treeID btrfsprim.ObjID, fn func(btrfsprim.Key) int) []btrfsprim.Key {
-	if !ts.AddTree(ctx, treeID) {
-		return nil
-	}
-	kvs := ts.trees[treeID].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.
-//
-// It is not nescessary to call AddTree for the tree first;
-// LeafToRoots will call it for you.
-func (ts *RebuiltTrees) LeafToRoots(ctx context.Context, treeID btrfsprim.ObjID, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
-	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) {
-			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
-	}
-	return ret
-}
-
-// Keys returns a map of all keys in node that would be valid in this tree.
-//
-// It is invalid (panic) to call Keys for a tree without having called
-// AddTree first.
-func (ts *RebuiltTrees) Keys(treeID btrfsprim.ObjID) *containers.SortedMap[btrfsprim.Key, keyio.ItemPtr] {
-	return &ts.trees[treeID].keys
-}
-
-// COWDistance returns how many COW-snapshots down from the 'child'
-// tree is from the 'parent' tree.
-//
-// It is invalid (panic) to call COWDistance for a tree without having
-// called AddTree for the child first.
-func (ts *RebuiltTrees) COWDistance(ctx context.Context, childID, parentID btrfsprim.ObjID) (dist int, ok bool) {
-	return ts.trees[childID].cowDistance(parentID)
-}
-
-// ListRoots returns a listing of all initialized trees and their root
-// nodes.
-//
-// Do not mutate the set of roots for a tree; it is a pointer to the
-// RebuiltTrees' internal set!
-func (ts *RebuiltTrees) 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
-	}
-	return ret
-}
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-2-g168b


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 +++++++++++++++------
 2 files changed, 111 insertions(+), 46 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

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
-}
-- 
cgit v1.2.3-2-g168b


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 ++--------------------
 1 file changed, 3 insertions(+), 31 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

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] {
-- 
cgit v1.2.3-2-g168b


From 0e73803c2f951fe9f9f0c09e8c2bf14ac0cd8ef2 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sat, 31 Dec 2022 10:13:00 -0700
Subject: rebuildnodes/btrees: Tune cache sizes

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

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index c0b7698..ef6da6f 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -85,8 +85,8 @@ func NewRebuiltForrest(
 		cbLookupUUID: cbLookupUUID,
 
 		trees:    make(map[btrfsprim.ObjID]*RebuiltTree),
-		allItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(16)),
-		incItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(16)),
+		allItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
+		incItems: containers.NewLRUCache[btrfsprim.ObjID, *itemIndex](textui.Tunable(8)),
 	}
 }
 
-- 
cgit v1.2.3-2-g168b


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 ++++++
 2 files changed, 18 insertions(+), 11 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

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]),
-- 
cgit v1.2.3-2-g168b


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 ++++++++++++----------
 2 files changed, 39 insertions(+), 33 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

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))
-- 
cgit v1.2.3-2-g168b


From 7ba1804756b43359a8b85e601b3836c8c4aac6cf Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Mon, 2 Jan 2023 18:26:51 -0700
Subject: rebuildnodes/btrees: Fix logging of the add-tree stack

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

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index de21d9a..6bfa297 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -150,7 +150,7 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			if !ok {
 				return false
 			}
-			if !ts.addTree(ctx, parentID, append(stack, treeID)) {
+			if !ts.addTree(ctx, parentID, stack) {
 				return false
 			}
 			tree.Parent, _ = ts.trees.Load(parentID)
-- 
cgit v1.2.3-2-g168b


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')

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


From 39dfaa5af4b628eee55b4aa583abf72ef5833f32 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Mon, 2 Jan 2023 18:52:27 -0700
Subject: rebuildnodes/btrees: Enhance logging around failure to add a tree

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

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 6f0ac96..3ed27b6 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -115,12 +115,13 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			ts.trees.Store(treeID, nil)
 		}
 	}()
-	if slices.Contains(treeID, stack) {
-		return false
-	}
 	stack = append(stack, treeID)
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree", stack)
 	dlog.Info(ctx, "adding tree...")
+	if slices.Contains(treeID, stack[:len(stack)-1]) {
+		dlog.Errorf(ctx, "failed to add tree: loop detected: %v", stack)
+		return false
+	}
 
 	tree := &RebuiltTree{
 		ID:      treeID,
@@ -140,10 +141,12 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 		root = ts.sb.BlockGroupRoot
 	default:
 		if !ts.addTree(ctx, btrfsprim.ROOT_TREE_OBJECTID, stack) {
+			dlog.Error(ctx, "failed to add tree: add ROOT_TREE")
 			return false
 		}
 		rootOff, rootItem, ok := ts.cbLookupRoot(ctx, treeID)
 		if !ok {
+			dlog.Error(ctx, "failed to add tree: lookup ROOT_ITEM")
 			return false
 		}
 		root = rootItem.ByteNr
@@ -155,9 +158,11 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			}
 			parentID, ok := ts.cbLookupUUID(ctx, rootItem.ParentUUID)
 			if !ok {
+				dlog.Error(ctx, "failed to add tree: lookup UUID")
 				return false
 			}
 			if !ts.addTree(ctx, parentID, stack) {
+				dlog.Error(ctx, "failed to add tree: add parent tree")
 				return false
 			}
 			tree.Parent, _ = ts.trees.Load(parentID)
-- 
cgit v1.2.3-2-g168b


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')

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-2-g168b


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')

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-2-g168b


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')

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-2-g168b


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 +++++-----
 2 files changed, 7 insertions(+), 7 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees')

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)
 		})
 }
 
-- 
cgit v1.2.3-2-g168b


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')

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-2-g168b


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')

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-2-g168b