containers: Rethink the caching libraries

4 files changed, 103 insertions, 97 deletions

diff --git a/lib/btrfsutil/open.go b/lib/btrfsutil/open.go
index c5ee314..abbd466 100644
--- a/lib/btrfsutil/open.go
+++ b/lib/btrfsutil/open.go
@@ -30,6 +30,7 @@ func Open(ctx context.Context, flag int, filenames ...string) (*btrfs.FS, error)
 			File: osFile,
 		}
 		bufFile := diskio.NewBufferedFile[btrfsvol.PhysicalAddr](
+			ctx,
 			typedFile,
 			//nolint:gomnd // False positive: gomnd.ignored-functions=[textui.Tunable] doesn't support type params.
 			textui.Tunable[btrfsvol.PhysicalAddr](16*1024), // block size: 16KiB
diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index b5c646d..811e1ac 100644
--- a/lib/btrfsutil/rebuilt_forrest.go
+++ b/lib/btrfsutil/rebuilt_forrest.go
@@ -7,7 +7,6 @@ package btrfsutil
 import (
 	"context"
 	"fmt"
-	"sync"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -140,12 +139,11 @@ type RebuiltForrest struct {
 
 	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]
+	leafs    containers.Cache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]
+	incItems containers.Cache[btrfsprim.ObjID, itemIndex]
+	excItems containers.Cache[btrfsprim.ObjID, itemIndex]
 
-	nodesMu sync.Mutex
-	nodes   containers.ARCache[btrfsvol.LogicalAddr, *btrfstree.Node]
+	nodes containers.Cache[btrfsvol.LogicalAddr, btrfstree.Node]
 }
 
 // NewRebuiltForrest returns a new RebuiltForrest instance.  The
@@ -158,23 +156,24 @@ func NewRebuiltForrest(file diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Supe
 		cb:    cb,
 
 		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
-		leafs: containers.ARCache[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]]{
-			MaxLen: textui.Tunable(8),
-		},
-		incItems: containers.ARCache[btrfsprim.ObjID, *itemIndex]{
-			MaxLen: textui.Tunable(8),
-		},
-		excItems: containers.ARCache[btrfsprim.ObjID, *itemIndex]{
-			MaxLen: textui.Tunable(8),
-		},
-
-		nodes: containers.ARCache[btrfsvol.LogicalAddr, *btrfstree.Node]{
-			MaxLen: textui.Tunable(8),
-			OnRemove: func(_ btrfsvol.LogicalAddr, node *btrfstree.Node) {
-				node.Free()
-			},
-		},
 	}
+
+	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 = ret.trees[treeID].uncachedLeafToRoots(ctx)
+			}))
+	ret.incItems = containers.NewARCache[btrfsprim.ObjID, itemIndex](textui.Tunable(8),
+		containers.SourceFunc[btrfsprim.ObjID, itemIndex](func(ctx context.Context, treeID btrfsprim.ObjID, incItems *itemIndex) {
+			*incItems = ret.trees[treeID].uncachedIncItems(ctx)
+		}))
+	ret.excItems = containers.NewARCache[btrfsprim.ObjID, itemIndex](textui.Tunable(8),
+		containers.SourceFunc[btrfsprim.ObjID, itemIndex](func(ctx context.Context, treeID btrfsprim.ObjID, excItems *itemIndex) {
+			*excItems = ret.trees[treeID].uncachedExcItems(ctx)
+		}))
+	ret.nodes = containers.NewARCache[btrfsvol.LogicalAddr, btrfstree.Node](textui.Tunable(8),
+		containers.SourceFunc[btrfsvol.LogicalAddr, btrfstree.Node](ret.readNode))
+
 	if ret.cb == nil {
 		ret.cb = noopRebuiltForrestCallbacks{
 			forrest: ret,
diff --git a/lib/btrfsutil/rebuilt_readitem.go b/lib/btrfsutil/rebuilt_readitem.go
index 03a7cdc..d3a2253 100644
--- a/lib/btrfsutil/rebuilt_readitem.go
+++ b/lib/btrfsutil/rebuilt_readitem.go
@@ -26,18 +26,14 @@ func (ptr ItemPtr) String() string {
 	return fmt.Sprintf("node@%v[%v]", ptr.Node, ptr.Slot)
 }
 
-func (ts *RebuiltForrest) readNode(ctx context.Context, laddr btrfsvol.LogicalAddr) *btrfstree.Node {
-	if cached, ok := ts.nodes.Load(laddr); ok {
-		dlog.Tracef(ctx, "cache-hit node@%v", laddr)
-		return cached
-	}
+func (ts *RebuiltForrest) readNode(ctx context.Context, laddr btrfsvol.LogicalAddr, out *btrfstree.Node) {
+	dlog.Debugf(ctx, "cache-miss node@%v, reading...", laddr)
 
 	graphInfo, ok := ts.graph.Nodes[laddr]
 	if !ok {
 		panic(fmt.Errorf("should not happen: node@%v is not mentioned in the in-memory graph", laddr))
 	}
 
-	dlog.Debugf(ctx, "cache-miss node@%v, reading...", laddr)
 	node, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](ts.file, ts.sb, laddr, btrfstree.NodeExpectations{
 		LAddr:      containers.OptionalValue(laddr),
 		Level:      containers.OptionalValue(graphInfo.Level),
@@ -56,22 +52,20 @@ func (ts *RebuiltForrest) readNode(ctx context.Context, laddr btrfsvol.LogicalAd
 	if err != nil {
 		panic(fmt.Errorf("should not happen: i/o error: %w", err))
 	}
-
-	ts.nodes.Store(laddr, node)
-
-	return node
+	*out = *node
 }
 
 func (ts *RebuiltForrest) readItem(ctx context.Context, ptr ItemPtr) btrfsitem.Item {
-	ts.nodesMu.Lock()
-	defer ts.nodesMu.Unlock()
 	if ts.graph.Nodes[ptr.Node].Level != 0 {
 		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for non-leaf node@%v", ptr.Node))
 	}
 	if ptr.Slot < 0 {
 		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for negative item slot: %v", ptr.Slot))
 	}
-	items := ts.readNode(ctx, ptr.Node).BodyLeaf
+
+	items := ts.nodes.Acquire(ctx, ptr.Node).BodyLeaf
+	defer ts.nodes.Release(ptr.Node)
+
 	if ptr.Slot >= len(items) {
 		panic(fmt.Errorf("should not happen: btrfsutil.RebuiltForrest.readItem called for out-of-bounds item slot: slot=%v len=%v",
 			ptr.Slot, len(items)))
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index 96d5a75..e65a3f6 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -47,33 +47,37 @@ type RebuiltTree struct {
 // 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 containers.LoadOrElse[btrfsprim.ObjID, map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr]](&tree.forrest.leafs, tree.ID, func(btrfsprim.ObjID) map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr] {
-		ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-nodes", fmt.Sprintf("tree=%v", tree.ID))
+	ret := *tree.forrest.leafs.Acquire(ctx, tree.ID)
+	tree.forrest.leafs.Release(tree.ID)
+	return ret
+}
 
-		nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
+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))
 
-		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)
-		}
+	nodeToRoots := make(map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr])
 
-		progress()
-		for _, node := range maps.SortedKeys(tree.forrest.graph.Nodes) {
-			tree.indexNode(ctx, node, nodeToRoots, progress, nil)
-		}
-		progressWriter.Done()
+	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)
+	}
 
-		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
-			}
+	progress()
+	for _, node := range maps.SortedKeys(tree.forrest.graph.Nodes) {
+		tree.indexNode(ctx, node, nodeToRoots, progress, nil)
+	}
+	progressWriter.Done()
+
+	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
-	})
+	}
+	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) {
@@ -136,8 +140,9 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 // 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] {
-	ctx = dlog.WithField(ctx, "btrfs.util.rebuilt-tree.index-inc-items", fmt.Sprintf("tree=%v", tree.ID))
-	return tree.items(ctx, &tree.forrest.incItems, tree.Roots.HasAny)
+	ret := *tree.forrest.incItems.Acquire(ctx, tree.ID)
+	tree.forrest.incItems.Release(tree.ID)
+	return ret
 }
 
 // RebuiltPotentialItems returns a map of items that could be added to
@@ -146,14 +151,25 @@ func (tree *RebuiltTree) RebuiltItems(ctx context.Context) *containers.SortedMap
 // 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)
+	tree.forrest.excItems.Release(tree.ID)
+	return ret
+}
+
+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)
+}
+
+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, &tree.forrest.excItems,
+	return tree.items(ctx,
 		func(roots containers.Set[btrfsvol.LogicalAddr]) bool {
 			return !tree.Roots.HasAny(roots)
 		})
 }
 
-type itemIndex = containers.SortedMap[btrfsprim.Key, ItemPtr]
+type itemIndex = *containers.SortedMap[btrfsprim.Key, ItemPtr]
 
 type itemStats struct {
 	Leafs    textui.Portion[int]
@@ -166,54 +182,50 @@ func (s itemStats) String() string {
 		s.Leafs, s.NumItems, s.NumDups)
 }
 
-func (tree *RebuiltTree) items(ctx context.Context, cache containers.Map[btrfsprim.ObjID, *itemIndex],
-	leafFn func(roots containers.Set[btrfsvol.LogicalAddr]) bool,
-) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+func (tree *RebuiltTree) items(ctx context.Context, leafFn func(roots containers.Set[btrfsvol.LogicalAddr]) bool) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
 
-	return containers.LoadOrElse(cache, tree.ID, func(btrfsprim.ObjID) *itemIndex {
-		var leafs []btrfsvol.LogicalAddr
-		for leaf, roots := range tree.leafToRoots(ctx) {
-			if leafFn(roots) {
-				leafs = append(leafs, leaf)
-			}
+	var leafs []btrfsvol.LogicalAddr
+	for leaf, roots := range tree.leafToRoots(ctx) {
+		if leafFn(roots) {
+			leafs = append(leafs, leaf)
 		}
-		slices.Sort(leafs)
+	}
+	slices.Sort(leafs)
 
-		var stats itemStats
-		stats.Leafs.D = len(leafs)
-		progressWriter := textui.NewProgress[itemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	var stats itemStats
+	stats.Leafs.D = len(leafs)
+	progressWriter := textui.NewProgress[itemStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
 
-		index := new(containers.SortedMap[btrfsprim.Key, ItemPtr])
-		for i, leaf := range leafs {
-			stats.Leafs.N = i
-			progressWriter.Set(stats)
-			for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
-				newPtr := ItemPtr{
-					Node: leaf,
-					Slot: j,
-				}
-				if oldPtr, exists := index.Load(itemKey); !exists {
+	index := new(containers.SortedMap[btrfsprim.Key, ItemPtr])
+	for i, leaf := range leafs {
+		stats.Leafs.N = i
+		progressWriter.Set(stats)
+		for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+			newPtr := ItemPtr{
+				Node: leaf,
+				Slot: j,
+			}
+			if oldPtr, exists := index.Load(itemKey); !exists {
+				index.Store(itemKey, newPtr)
+				stats.NumItems++
+			} else {
+				if tree.RebuiltShouldReplace(oldPtr.Node, newPtr.Node) {
 					index.Store(itemKey, newPtr)
-					stats.NumItems++
-				} else {
-					if tree.RebuiltShouldReplace(oldPtr.Node, newPtr.Node) {
-						index.Store(itemKey, newPtr)
-					}
-					stats.NumDups++
 				}
-				progressWriter.Set(stats)
+				stats.NumDups++
 			}
-		}
-		if stats.Leafs.N > 0 {
-			stats.Leafs.N = len(leafs)
 			progressWriter.Set(stats)
-			progressWriter.Done()
 		}
+	}
+	if stats.Leafs.N > 0 {
+		stats.Leafs.N = len(leafs)
+		progressWriter.Set(stats)
+		progressWriter.Done()
+	}
 
-		return index
-	})
+	return index
 }
 
 // main public API /////////////////////////////////////////////////////////////////////////////////////////////////////