rebuildnodes: Parallelize I/O and CPU

2 files changed, 18 insertions, 11 deletions

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]),