Merge branch 'lukeshu/rebuild-nodes-take5'

3 files changed, 93 insertions, 46 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 3eeea7f..9ec2849 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -7,7 +7,6 @@ package btrees
 import (
 	"context"
 
-	"git.lukeshu.com/go/typedsync"
 	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
@@ -21,6 +20,12 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
+type Callbacks interface {
+	AddedItem(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
+	LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool)
+	LookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
+}
+
 // RebuiltForrest is an abstraction for rebuilding and accessing
 // potentially broken btrees.
 //
@@ -56,14 +61,11 @@ type RebuiltForrest struct {
 	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)
+	cb    Callbacks
 
 	// mutable
-	trees    typedsync.Map[btrfsprim.ObjID, *RebuiltTree]
+	treesMu  nestedMutex
+	trees    map[btrfsprim.ObjID]*RebuiltTree // must hold .treesMu to access
 	leafs    containers.ARCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]
 	incItems containers.ARCache[btrfsprim.ObjID, *itemIndex]
 	excItems containers.ARCache[btrfsprim.ObjID, *itemIndex]
@@ -71,21 +73,14 @@ type RebuiltForrest struct {
 
 // 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 {
+func NewRebuiltForrest(sb btrfstree.Superblock, graph pkggraph.Graph, keyIO *keyio.Handle, cb Callbacks) *RebuiltForrest {
 	return &RebuiltForrest{
 		sb:    sb,
 		graph: graph,
 		keyIO: keyIO,
+		cb:    cb,
 
-		cbAddedItem:  cbAddedItem,
-		cbLookupRoot: cbLookupRoot,
-		cbLookupUUID: cbLookupUUID,
-
+		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
 		leafs: containers.ARCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]{
 			MaxLen: textui.Tunable(8),
 		},
@@ -104,22 +99,23 @@ func NewRebuiltForrest(
 //
 // The tree is initialized with the normal root node of the tree.
 func (ts *RebuiltForrest) Tree(ctx context.Context, treeID btrfsprim.ObjID) *RebuiltTree {
+	ctx = ts.treesMu.Lock(ctx)
+	defer ts.treesMu.Unlock()
 	if !ts.addTree(ctx, treeID, nil) {
 		return nil
 	}
-	tree, _ := ts.trees.Load(treeID)
-	return tree
+	return ts.trees[treeID]
 }
 
 func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, stack []btrfsprim.ObjID) (ok bool) {
-	if tree, ok := ts.trees.Load(treeID); ok {
+	if tree, ok := ts.trees[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)
+			// trees will call .flushNegativeCache().
+			ts.trees[treeID] = nil
 		}
 	}()
 	stack = append(stack, treeID)
@@ -150,7 +146,7 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			dlog.Error(ctx, "failed to add tree: add ROOT_TREE")
 			return false
 		}
-		rootOff, rootItem, ok := ts.cbLookupRoot(ctx, treeID)
+		rootOff, rootItem, ok := ts.cb.LookupRoot(ctx, treeID)
 		if !ok {
 			dlog.Error(ctx, "failed to add tree: lookup ROOT_ITEM")
 			return false
@@ -162,7 +158,7 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 			if !ts.addTree(ctx, btrfsprim.UUID_TREE_OBJECTID, stack) {
 				return false
 			}
-			parentID, ok := ts.cbLookupUUID(ctx, rootItem.ParentUUID)
+			parentID, ok := ts.cb.LookupUUID(ctx, rootItem.ParentUUID)
 			if !ok {
 				dlog.Error(ctx, "failed to add tree: lookup UUID")
 				return false
@@ -171,11 +167,11 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 				dlog.Error(ctx, "failed to add tree: add parent tree")
 				return false
 			}
-			tree.Parent, _ = ts.trees.Load(parentID)
+			tree.Parent = ts.trees[parentID]
 		}
 	}
 
-	ts.trees.Store(treeID, tree)
+	ts.trees[treeID] = tree
 	if root != 0 {
 		tree.AddRoot(ctx, root)
 	}
@@ -183,18 +179,29 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 	return true
 }
 
+func (ts *RebuiltForrest) flushNegativeCache(ctx context.Context) {
+	_ = ts.treesMu.Lock(ctx)
+	defer ts.treesMu.Unlock()
+	for treeID, tree := range ts.trees {
+		if tree == nil {
+			delete(ts.trees, treeID)
+		}
+	}
+}
+
 // 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] {
+func (ts *RebuiltForrest) ListRoots(ctx context.Context) map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
+	_ = ts.treesMu.Lock(ctx)
+	defer ts.treesMu.Unlock()
 	ret := make(map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr])
-	ts.trees.Range(func(treeID btrfsprim.ObjID, tree *RebuiltTree) bool {
+	for treeID, tree := range ts.trees {
 		if tree != nil {
 			ret[treeID] = tree.Roots
 		}
-		return true
-	})
+	}
 	return ret
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/nestedlock.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/nestedlock.go
new file mode 100644
index 0000000..c1ffa18
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/nestedlock.go
@@ -0,0 +1,45 @@
+// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrees
+
+import (
+	"context"
+	"sync"
+)
+
+// A nestedMutex is like a sync.Mutex, but while it is locked by call
+// 'A', may be simultaneously locked by subsequent calls if the
+// subsequent calls use a Context descended from the one returned by
+// the 'A' call to .Lock().
+type nestedMutex struct {
+	inner sync.Mutex
+	depth int
+}
+
+type nestedMutexCtxKey struct{}
+
+// Lock locks the mutex.  It is invalid to use a Context returned from
+// Lock in a different goroutine than the one it was created in.  It
+// is invalid to use a Context returned from Lock after the mutex has
+// subsequently become unlocked.
+func (m *nestedMutex) Lock(ctx context.Context) context.Context {
+	if other, ok := ctx.Value(nestedMutexCtxKey{}).(*nestedMutex); ok && other == m {
+		m.depth++
+		return ctx
+	}
+	m.inner.Lock()
+	return context.WithValue(ctx, nestedMutexCtxKey{}, m)
+}
+
+// Unlock unlocks the mutex.  It is invalid to call Unlock if the
+// mutex is not already locked.  It is invalid to call Unlock from
+// multiple goroutines simultaneously.
+func (m *nestedMutex) Unlock() {
+	if m.depth > 0 {
+		m.depth--
+	} else {
+		m.inner.Unlock()
+	}
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index c9d0fa4..eab3eb2 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -15,7 +15,6 @@ import (
 	"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"
@@ -99,10 +98,10 @@ func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAdd
 	// 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 {
+	for _, kp := range tree.forrest.graph.EdgesTo[node] {
+		if !tree.isOwnerOK(tree.forrest.graph.Nodes[kp.FromNode].Owner, tree.forrest.graph.Nodes[kp.FromNode].Generation) {
+			continue
+		}
 		tree.indexNode(ctx, kp.FromNode, index, progress, stack)
 		if len(index[kp.FromNode]) > 0 {
 			if roots == nil {
@@ -200,7 +199,7 @@ func (tree *RebuiltTree) items(ctx context.Context, cache containers.Map[btrfspr
 					index.Store(itemKey, newPtr)
 					stats.NumItems++
 				} else {
-					if tree.shouldReplace(oldPtr.Node, newPtr.Node) {
+					if tree.ShouldReplace(oldPtr.Node, newPtr.Node) {
 						index.Store(itemKey, newPtr)
 					}
 					stats.NumDups++
@@ -218,7 +217,7 @@ func (tree *RebuiltTree) items(ctx context.Context, cache containers.Map[btrfspr
 	})
 }
 
-func (tree *RebuiltTree) shouldReplace(oldNode, newNode btrfsvol.LogicalAddr) bool {
+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 {
@@ -270,6 +269,7 @@ 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))
+	dlog.Info(ctx, "adding root...")
 
 	leafToRoots := tree.leafToRoots(ctx)
 
@@ -288,7 +288,7 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 		progressWriter.Set(stats)
 
 		for _, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
-			tree.forrest.cbAddedItem(ctx, tree.ID, itemKey)
+			tree.forrest.cb.AddedItem(ctx, tree.ID, itemKey)
 			stats.AddedItems++
 			progressWriter.Set(stats)
 		}
@@ -302,12 +302,7 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	tree.forrest.excItems.Delete(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 {
-			if otherTree == nil {
-				tree.forrest.trees.Delete(otherTreeID)
-			}
-			return true
-		})
+		tree.forrest.flushNegativeCache(ctx)
 	}
 }
 
@@ -329,10 +324,10 @@ func (tree *RebuiltTree) COWDistance(parentID btrfsprim.ObjID) (dist int, ok boo
 }
 
 // ReadItem reads an item from a tree.
-func (tree *RebuiltTree) ReadItem(ctx context.Context, key btrfsprim.Key) (item btrfsitem.Item, ok bool) {
+func (tree *RebuiltTree) ReadItem(ctx context.Context, key btrfsprim.Key) btrfsitem.Item {
 	ptr, ok := tree.Items(ctx).Load(key)
 	if !ok {
-		return nil, false
+		panic(fmt.Errorf("should not happen: btrees.RebuiltTree.ReadItem called for not-included key: %v", key))
 	}
 	return tree.forrest.keyIO.ReadItem(ctx, ptr)
 }