Merge branch 'lukeshu/misc'

1 files changed, 88 insertions, 46 deletions

diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index b996bdb..ffb2e5f 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -37,21 +37,52 @@ type RebuiltTree struct {
 	// derived from tree.Roots, which is why it's OK if they get
 	// evicted.
 	//
-	//  1. tree.leafToRoots()           = tree.forrest.leafs.Acquire(tree.ID)
-	//  2. tree.RebuiltItems()          = tree.forrest.incItems.Acquire(tree.ID)
-	//  3. tree.RebuiltPotentialItems() = tree.forrest.excItems.Acquire(tree.ID)
+	//  1. tree.acquireLeafToRoots()           = tree.forrest.leafs.Acquire(tree.ID)
+	//  2. tree.RebuiltAcquireItems()          = tree.forrest.incItems.Acquire(tree.ID)
+	//  3. tree.RebuiltAcquirePotentialItems() = tree.forrest.excItems.Acquire(tree.ID)
 }
 
-// evictable member 1: .leafToRoots() //////////////////////////////////////////////////////////////////////////////////
+type rebuiltSharedCache struct {
+	leafs    containers.Cache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]
+	incItems containers.Cache[btrfsprim.ObjID, containers.SortedMap[btrfsprim.Key, ItemPtr]]
+	excItems containers.Cache[btrfsprim.ObjID, containers.SortedMap[btrfsprim.Key, ItemPtr]]
+}
 
-// 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] {
-	ret := *tree.forrest.leafs.Acquire(ctx, tree.ID)
-	tree.forrest.leafs.Release(tree.ID)
+func makeRebuiltSharedCache(forrest *RebuiltForrest) rebuiltSharedCache {
+	var ret rebuiltSharedCache
+	ret.leafs = containers.NewARCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]](
+		textui.Tunable(8),
+		containers.SourceFunc[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]](
+			func(ctx context.Context, treeID btrfsprim.ObjID, leafs *map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]) {
+				*leafs = forrest.trees[treeID].uncachedLeafToRoots(ctx)
+			}))
+	ret.incItems = containers.NewARCache[btrfsprim.ObjID, containers.SortedMap[btrfsprim.Key, ItemPtr]](
+		textui.Tunable(8),
+		containers.SourceFunc[btrfsprim.ObjID, containers.SortedMap[btrfsprim.Key, ItemPtr]](
+			func(ctx context.Context, treeID btrfsprim.ObjID, incItems *containers.SortedMap[btrfsprim.Key, ItemPtr]) {
+				*incItems = forrest.trees[treeID].uncachedIncItems(ctx)
+			}))
+	ret.excItems = containers.NewARCache[btrfsprim.ObjID, containers.SortedMap[btrfsprim.Key, ItemPtr]](
+		textui.Tunable(8),
+		containers.SourceFunc[btrfsprim.ObjID, containers.SortedMap[btrfsprim.Key, ItemPtr]](
+			func(ctx context.Context, treeID btrfsprim.ObjID, excItems *containers.SortedMap[btrfsprim.Key, ItemPtr]) {
+				*excItems = forrest.trees[treeID].uncachedExcItems(ctx)
+			}))
 	return ret
 }
 
+// evictable member 1: .acquireLeafToRoots() ///////////////////////////////////////////////////////////////////////////
+
+// acquireLeafToRoots returns all leafs (lvl=0) in the filesystem that
+// pass .isOwnerOK, whether or not they're in the tree.
+func (tree *RebuiltTree) acquireLeafToRoots(ctx context.Context) map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr] {
+	return *tree.forrest.leafs.Acquire(ctx, tree.ID)
+}
+
+func (tree *RebuiltTree) releaseLeafToRoots() {
+	tree.forrest.leafs.Release(tree.ID)
+}
+
 func (tree *RebuiltTree) uncachedLeafToRoots(ctx context.Context) map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr] {
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-nodes", fmt.Sprintf("tree=%v", tree.ID))
 
@@ -85,7 +116,7 @@ func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAdd
 	if err := ctx.Err(); err != nil {
 		return
 	}
-	if _, done := index[node]; done {
+	if maps.HasKey(index, node) {
 		return
 	}
 	if slices.Contains(node, stack) {
@@ -133,72 +164,81 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 	}
 }
 
-// evictable members 2 and 3: .RebuiltItems() and .RebuiltPotentialItems() /////////////////////////////////////////////
+// evictable members 2 and 3: .Rebuilt{Acquire,Release}{Potential,}Items() /////////////////////////////////////////////
 
-// RebuiltItems returns a map of the items contained in this tree.
+// RebuiltAcquireItems 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) RebuiltItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
-	ret := *tree.forrest.incItems.Acquire(ctx, tree.ID)
+//
+// When done with the map, call .RebuiltReleaseItems().
+func (tree *RebuiltTree) RebuiltAcquireItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+	return tree.forrest.incItems.Acquire(ctx, tree.ID)
+}
+
+// RebuiltReleaseItems releases resources after a call to
+// .RebuiltAcquireItems().
+func (tree *RebuiltTree) RebuiltReleaseItems() {
 	tree.forrest.incItems.Release(tree.ID)
-	return ret
 }
 
-// RebuiltPotentialItems returns a map of items that could be added to
-// this tree with .RebuiltAddRoot().
+// RebuiltAcquirePotentialItems returns a map of items that could be
+// added to this tree with .RebuiltAddRoot().
 //
 // Do not mutate the returned map; it is a pointer to the
 // RebuiltTree's internal map!
-func (tree *RebuiltTree) RebuiltPotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
-	ret := *tree.forrest.excItems.Acquire(ctx, tree.ID)
+//
+// When done with the map, call .RebuiltReleasePotentialItems().
+func (tree *RebuiltTree) RebuiltAcquirePotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+	return tree.forrest.excItems.Acquire(ctx, tree.ID)
+}
+
+// RebuiltReleasePotentialItems releases resources after a call to
+// .RebuiltAcquirePotentialItems().
+func (tree *RebuiltTree) RebuiltReleasePotentialItems() {
 	tree.forrest.excItems.Release(tree.ID)
-	return ret
 }
 
-func (tree *RebuiltTree) uncachedIncItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+func (tree *RebuiltTree) uncachedIncItems(ctx context.Context) containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-inc-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, tree.Roots.HasAny)
+	return tree.items(ctx, true)
 }
 
-func (tree *RebuiltTree) uncachedExcItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+func (tree *RebuiltTree) uncachedExcItems(ctx context.Context) containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-exc-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx,
-		func(roots containers.Set[btrfsvol.LogicalAddr]) bool {
-			return !tree.Roots.HasAny(roots)
-		})
+	return tree.items(ctx, false)
 }
 
-type itemIndex = *containers.SortedMap[btrfsprim.Key, ItemPtr]
-
-type itemStats struct {
+type rebuiltItemStats struct {
 	Leafs    textui.Portion[int]
 	NumItems int
 	NumDups  int
 }
 
-func (s itemStats) String() string {
+func (s rebuiltItemStats) String() string {
 	return textui.Sprintf("%v (%v items, %v dups)",
 		s.Leafs, s.NumItems, s.NumDups)
 }
 
-func (tree *RebuiltTree) items(ctx context.Context, leafFn func(roots containers.Set[btrfsvol.LogicalAddr]) bool) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+func (tree *RebuiltTree) items(ctx context.Context, inc bool) containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
 
 	var leafs []btrfsvol.LogicalAddr
-	for leaf, roots := range tree.leafToRoots(ctx) {
-		if leafFn(roots) {
+	for leaf, roots := range tree.acquireLeafToRoots(ctx) {
+		if tree.Roots.HasAny(roots) == inc {
 			leafs = append(leafs, leaf)
 		}
 	}
+	tree.releaseLeafToRoots()
 	slices.Sort(leafs)
 
-	var stats itemStats
+	var stats rebuiltItemStats
 	stats.Leafs.D = len(leafs)
-	progressWriter := textui.NewProgress[itemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progressWriter := textui.NewProgress[rebuiltItemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 
-	index := new(containers.SortedMap[btrfsprim.Key, ItemPtr])
+	var index containers.SortedMap[btrfsprim.Key, ItemPtr]
 	for i, leaf := range leafs {
 		stats.Leafs.N = i
 		progressWriter.Set(stats)
@@ -220,7 +260,7 @@ func (tree *RebuiltTree) items(ctx context.Context, leafFn func(roots containers
 		}
 	}
 	if stats.Leafs.N > 0 {
-		stats.Leafs.N = len(leafs)
+		stats.Leafs.N = stats.Leafs.D
 		progressWriter.Set(stats)
 		progressWriter.Done()
 	}
@@ -262,13 +302,13 @@ func (tree *RebuiltTree) RebuiltShouldReplace(oldNode, newNode btrfsvol.LogicalA
 	}
 }
 
-type rootStats struct {
+type rebuiltRootStats struct {
 	Leafs      textui.Portion[int]
 	AddedLeafs int
 	AddedItems int
 }
 
-func (s rootStats) String() string {
+func (s rebuiltRootStats) String() string {
 	return textui.Sprintf("%v (added %v leafs, added %v items)",
 		s.Leafs, s.AddedLeafs, s.AddedItems)
 }
@@ -282,11 +322,10 @@ func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.L
 	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.add-root", fmt.Sprintf("tree=%v rootNode=%v", tree.ID, rootNode))
 	dlog.Info(ctx, "adding root...")
 
-	leafToRoots := tree.leafToRoots(ctx)
-
-	var stats rootStats
+	var stats rebuiltRootStats
+	leafToRoots := tree.acquireLeafToRoots(ctx)
 	stats.Leafs.D = len(leafToRoots)
-	progressWriter := textui.NewProgress[rootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	progressWriter := textui.NewProgress[rebuiltRootStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 	for i, leaf := range maps.SortedKeys(leafToRoots) {
 		stats.Leafs.N = i
 		progressWriter.Set(stats)
@@ -305,6 +344,7 @@ func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.L
 		}
 	}
 	stats.Leafs.N = len(leafToRoots)
+	tree.releaseLeafToRoots()
 	progressWriter.Set(stats)
 	progressWriter.Done()
 
@@ -335,7 +375,8 @@ func (tree *RebuiltTree) RebuiltCOWDistance(parentID btrfsprim.ObjID) (dist int,
 
 // ReadItem reads an item from a tree.
 func (tree *RebuiltTree) ReadItem(ctx context.Context, key btrfsprim.Key) btrfsitem.Item {
-	ptr, ok := tree.RebuiltItems(ctx).Load(key)
+	ptr, ok := tree.RebuiltAcquireItems(ctx).Load(key)
+	tree.RebuiltReleaseItems()
 	if !ok {
 		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltTree.ReadItem called for not-included key: %v", key))
 	}
@@ -352,13 +393,14 @@ func (tree *RebuiltTree) RebuiltLeafToRoots(ctx context.Context, leaf btrfsvol.L
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
 	ret := make(containers.Set[btrfsvol.LogicalAddr])
-	for root := range tree.leafToRoots(ctx)[leaf] {
+	for root := range tree.acquireLeafToRoots(ctx)[leaf] {
 		if tree.Roots.Has(root) {
 			panic(fmt.Errorf("should not happen: (tree=%v).RebuiltLeafToRoots(leaf=%v): tree contains root=%v but not leaf",
 				tree.ID, leaf, root))
 		}
 		ret.Insert(root)
 	}
+	tree.releaseLeafToRoots()
 	if len(ret) == 0 {
 		return nil
 	}