rebuildnodes/btrees: Have the concurrency story for .trees make more sense

4 files changed, 70 insertions, 21 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 4f4e05d..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"
@@ -65,7 +64,8 @@ type RebuiltForrest struct {
 	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]
@@ -80,6 +80,7 @@ func NewRebuiltForrest(sb btrfstree.Superblock, graph pkggraph.Graph, keyIO *key
 		keyIO: keyIO,
 		cb:    cb,
 
+		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
 		leafs: containers.ARCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]{
 			MaxLen: textui.Tunable(8),
 		},
@@ -98,22 +99,23 @@ func NewRebuiltForrest(sb btrfstree.Superblock, graph pkggraph.Graph, keyIO *key
 //
 // 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 call .flushNegativeCache().
-			ts.trees.Store(treeID, nil)
+			ts.trees[treeID] = nil
 		}
 	}()
 	stack = append(stack, treeID)
@@ -165,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)
 	}
@@ -177,13 +179,14 @@ func (ts *RebuiltForrest) addTree(ctx context.Context, treeID btrfsprim.ObjID, s
 	return true
 }
 
-func (ts *RebuiltForrest) flushNegativeCache() {
-	ts.trees.Range(func(treeID btrfsprim.ObjID, tree *RebuiltTree) bool {
+func (ts *RebuiltForrest) flushNegativeCache(ctx context.Context) {
+	_ = ts.treesMu.Lock(ctx)
+	defer ts.treesMu.Unlock()
+	for treeID, tree := range ts.trees {
 		if tree == nil {
-			ts.trees.Delete(treeID)
+			delete(ts.trees, treeID)
 		}
-		return true
-	})
+	}
 }
 
 // ListRoots returns a listing of all initialized trees and their root
@@ -191,13 +194,14 @@ func (ts *RebuiltForrest) flushNegativeCache() {
 //
 // 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 d64d7d2..9ddd3c5 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -303,7 +303,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.flushNegativeCache()
+		tree.forrest.flushNegativeCache(ctx)
 	}
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index bc4d7c9..dc78c2e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -71,7 +71,7 @@ type treeAugmentQueue struct {
 
 type Rebuilder interface {
 	Rebuild(context.Context) error
-	ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
+	ListRoots(context.Context) map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
 }
 
 func NewRebuilder(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (Rebuilder, error) {
@@ -90,8 +90,8 @@ func NewRebuilder(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	return o, nil
 }
 
-func (o *rebuilder) ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
-	return o.rebuilt.ListRoots()
+func (o *rebuilder) ListRoots(ctx context.Context) map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
+	return o.rebuilt.ListRoots(ctx)
 }
 
 func (o *rebuilder) Rebuild(ctx context.Context) error {