btrfsutil: RebuiltForrest: Add a lax-ancestor mode

4 files changed, 163 insertions, 39 deletions

diff --git a/lib/btrfs/btrfstree/path.go b/lib/btrfs/btrfstree/path.go
index acc0e73..327a39b 100644
--- a/lib/btrfs/btrfstree/path.go
+++ b/lib/btrfs/btrfstree/path.go
@@ -12,6 +12,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
 // Path is a path from the superblock or a ROOT_ITEM to a node or
@@ -137,6 +138,7 @@ func checkOwner(
 	ctx context.Context, forrest Forrest, treeID btrfsprim.ObjID,
 	ownerToCheck btrfsprim.ObjID, genToCheck btrfsprim.Generation,
 ) error {
+	var stack []btrfsprim.ObjID
 	for {
 		if ownerToCheck == treeID {
 			return nil
@@ -154,6 +156,12 @@ func checkOwner(
 				ownerToCheck, genToCheck, err)
 		}
 
+		stack = append(stack, treeID)
+		if slices.Contains(parentID, stack) {
+			// Don't get stuck in an infinite loop if there's a cycle.
+			parentID = 0
+		}
+
 		if parentID == 0 && parentGen == 0 {
 			return fmt.Errorf("owner=%v is not acceptable in this tree",
 				ownerToCheck)
diff --git a/lib/btrfsutil/rebuilt_forrest.go b/lib/btrfsutil/rebuilt_forrest.go
index 019d824..900e725 100644
--- a/lib/btrfsutil/rebuilt_forrest.go
+++ b/lib/btrfsutil/rebuilt_forrest.go
@@ -7,6 +7,7 @@ package btrfsutil
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -64,14 +65,18 @@ import (
 // NewRebuiltForrest().
 type RebuiltForrest struct {
 	// static
-	inner btrfs.ReadableFS
-	graph Graph
-	cb    RebuiltForrestCallbacks
+	inner        btrfs.ReadableFS
+	graph        Graph
+	cb           RebuiltForrestCallbacks
+	laxAncestors bool
 
 	// mutable
 
-	treesMu nestedMutex
-	trees   map[btrfsprim.ObjID]*RebuiltTree // must hold .treesMu to access
+	treesMu         nestedMutex
+	trees           map[btrfsprim.ObjID]*RebuiltTree // must hold .treesMu to access
+	commitTreesOnce sync.Once
+	treesCommitted  bool // must hold .treesMu to access
+	treesCommitter  btrfsprim.ObjID
 
 	rebuiltSharedCache
 }
@@ -81,11 +86,23 @@ type RebuiltForrest struct {
 // The `cb` RebuiltForrestCallbacks may be nil.  If `cb` also
 // implements RebuiltForrestExtendedCallbacks, then a series of
 // .AddedItem() calls will be made before each call to .AddedRoot().
-func NewRebuiltForrest(fs btrfs.ReadableFS, graph Graph, cb RebuiltForrestCallbacks) *RebuiltForrest {
+//
+// `laxAncestors` is whether or not an error instantiating an ancestor
+// tree should prevent instantiating an descendant tree (lax=false
+// prevents it, lax=true allows it).
+//
+//   - `laxAncestors` inhibits calls to
+//     RebuiltForrestExtendedCallbacks.AddedItem().
+//
+//   - `laxAncestors` causes a call to RebuiltTree.RebuiltAddRoot on
+//     the ROOT_TREE or the UUID_TREE to panic if a tree other than the
+//     ROOT_TREE or the UUID_TREE has been read from.
+func NewRebuiltForrest(fs btrfs.ReadableFS, graph Graph, cb RebuiltForrestCallbacks, laxAncestors bool) *RebuiltForrest {
 	ret := &RebuiltForrest{
-		inner: fs,
-		graph: graph,
-		cb:    cb,
+		inner:        fs,
+		graph:        graph,
+		cb:           cb,
+		laxAncestors: laxAncestors,
 
 		trees: make(map[btrfsprim.ObjID]*RebuiltTree),
 	}
@@ -100,6 +117,26 @@ func NewRebuiltForrest(fs btrfs.ReadableFS, graph Graph, cb RebuiltForrestCallba
 	return ret
 }
 
+func (ts *RebuiltForrest) commitTrees(ctx context.Context, treeID btrfsprim.ObjID) {
+	if treeID == btrfsprim.ROOT_TREE_OBJECTID || treeID == btrfsprim.UUID_TREE_OBJECTID {
+		return
+	}
+	ts.commitTreesOnce.Do(func() {
+		if !ts.laxAncestors {
+			return
+		}
+		ctx = ts.treesMu.Lock(ctx)
+		if !ts.treesCommitted {
+			// Make sure ROOT_TREE and UUID_TREE are ready for reading.
+			_, _ = ts.RebuiltTree(ctx, btrfsprim.ROOT_TREE_OBJECTID)
+			_, _ = ts.RebuiltTree(ctx, btrfsprim.UUID_TREE_OBJECTID)
+			ts.treesCommitted = true
+			ts.treesCommitter = treeID
+		}
+		ts.treesMu.Unlock()
+	})
+}
+
 // RebuiltTree returns a given tree, initializing it if nescessary.
 //
 // The tree is initialized with the normal root node of the tree.
@@ -111,7 +148,7 @@ func (ts *RebuiltForrest) RebuiltTree(ctx context.Context, treeID btrfsprim.ObjI
 	defer ts.treesMu.Unlock()
 	ts.rebuildTree(ctx, treeID, nil)
 	tree := ts.trees[treeID]
-	if tree.ancestorLoop && tree.rootErr == nil {
+	if tree.ancestorLoop && tree.rootErr == nil && tree.ancestorRoot == 0 {
 		var loop []btrfsprim.ObjID
 		for ancestor := tree; true; ancestor = ancestor.Parent {
 			loop = append(loop, ancestor.ID)
@@ -119,7 +156,11 @@ func (ts *RebuiltForrest) RebuiltTree(ctx context.Context, treeID btrfsprim.ObjI
 				break
 			}
 		}
-		tree.rootErr = fmt.Errorf("loop detected: %v", loop)
+		if ts.laxAncestors {
+			tree.ancestorRoot = loop[len(loop)-2]
+		} else {
+			tree.rootErr = fmt.Errorf("loop detected: %v", loop)
+		}
 	}
 	if tree.rootErr != nil {
 		return nil, tree.rootErr
@@ -182,7 +223,9 @@ func (ts *RebuiltForrest) rebuildTree(ctx context.Context, treeID btrfsprim.ObjI
 			ts.trees[treeID].ParentGen = rootOff
 			parentID, ok := ts.cb.LookupUUID(ctx, rootItem.ParentUUID)
 			if !ok {
-				ts.trees[treeID].rootErr = fmt.Errorf("failed to look up UUID: %v", rootItem.ParentUUID)
+				if !ts.laxAncestors {
+					ts.trees[treeID].rootErr = fmt.Errorf("failed to look up UUID: %v", rootItem.ParentUUID)
+				}
 				return
 			}
 			ts.rebuildTree(ctx, parentID, stack)
@@ -191,7 +234,7 @@ func (ts *RebuiltForrest) rebuildTree(ctx context.Context, treeID btrfsprim.ObjI
 			case ts.trees[treeID].Parent.ancestorLoop:
 				ts.trees[treeID].ancestorLoop = true
 				return
-			case ts.trees[treeID].Parent.rootErr != nil:
+			case !ts.laxAncestors && ts.trees[treeID].Parent.rootErr != nil:
 				ts.trees[treeID].rootErr = fmt.Errorf("failed to rebuild parent tree: %v: %w", parentID, ts.trees[treeID].Parent.rootErr)
 				return
 			}
diff --git a/lib/btrfsutil/rebuilt_forrest_test.go b/lib/btrfsutil/rebuilt_forrest_test.go
index 96749c4..8bbb50a 100644
--- a/lib/btrfsutil/rebuilt_forrest_test.go
+++ b/lib/btrfsutil/rebuilt_forrest_test.go
@@ -106,26 +106,60 @@ func TestRebuiltTreeCycles(t *testing.T) {
 		},
 	}
 
-	rfs := NewRebuiltForrest(nil, Graph{}, cbs)
-
-	tree, err := rfs.RebuiltTree(ctx, 306)
-	assert.EqualError(t, err, `loop detected: [306 305 304 303 305]`)
-	assert.Nil(t, tree)
-
-	assert.NotNil(t, rfs.trees[305])
-	tree, err = rfs.RebuiltTree(ctx, 305)
-	assert.EqualError(t, err, `loop detected: [305 304 303 305]`)
-	assert.Nil(t, tree)
-
-	assert.NotNil(t, rfs.trees[304])
-	tree, err = rfs.RebuiltTree(ctx, 304)
-	assert.EqualError(t, err, `loop detected: [304 303 305 304]`)
-	assert.Nil(t, tree)
-
-	assert.NotNil(t, rfs.trees[303])
-	tree, err = rfs.RebuiltTree(ctx, 303)
-	assert.EqualError(t, err, `loop detected: [303 305 304 303]`)
-	assert.Nil(t, tree)
+	t.Run("strict", func(t *testing.T) {
+		t.Parallel()
+		rfs := NewRebuiltForrest(nil, Graph{}, cbs, false)
+
+		tree, err := rfs.RebuiltTree(ctx, 306)
+		assert.EqualError(t, err, `loop detected: [306 305 304 303 305]`)
+		assert.Nil(t, tree)
+
+		assert.NotNil(t, rfs.trees[305])
+		tree, err = rfs.RebuiltTree(ctx, 305)
+		assert.EqualError(t, err, `loop detected: [305 304 303 305]`)
+		assert.Nil(t, tree)
+
+		assert.NotNil(t, rfs.trees[304])
+		tree, err = rfs.RebuiltTree(ctx, 304)
+		assert.EqualError(t, err, `loop detected: [304 303 305 304]`)
+		assert.Nil(t, tree)
+
+		assert.NotNil(t, rfs.trees[303])
+		tree, err = rfs.RebuiltTree(ctx, 303)
+		assert.EqualError(t, err, `loop detected: [303 305 304 303]`)
+		assert.Nil(t, tree)
+	})
+	t.Run("lax", func(t *testing.T) {
+		t.Parallel()
+		rfs := NewRebuiltForrest(nil, Graph{}, cbs, true)
+
+		tree, err := rfs.RebuiltTree(ctx, 306)
+		assert.NoError(t, err)
+		assert.NotNil(t, tree)
+		assert.True(t, tree.ancestorLoop)
+		assert.Equal(t, btrfsprim.ObjID(303), tree.ancestorRoot)
+
+		assert.NotNil(t, rfs.trees[305])
+		tree, err = rfs.RebuiltTree(ctx, 305)
+		assert.NoError(t, err)
+		assert.NotNil(t, tree)
+		assert.True(t, tree.ancestorLoop)
+		assert.Equal(t, btrfsprim.ObjID(303), tree.ancestorRoot)
+
+		assert.NotNil(t, rfs.trees[304])
+		tree, err = rfs.RebuiltTree(ctx, 304)
+		assert.NoError(t, err)
+		assert.NotNil(t, tree)
+		assert.True(t, tree.ancestorLoop)
+		assert.Equal(t, btrfsprim.ObjID(305), tree.ancestorRoot)
+
+		assert.NotNil(t, rfs.trees[303])
+		tree, err = rfs.RebuiltTree(ctx, 303)
+		assert.NoError(t, err)
+		assert.NotNil(t, tree)
+		assert.True(t, tree.ancestorLoop)
+		assert.Equal(t, btrfsprim.ObjID(304), tree.ancestorRoot)
+	})
 }
 
 func TestRebuiltTreeParentErr(t *testing.T) {
@@ -185,9 +219,21 @@ func TestRebuiltTreeParentErr(t *testing.T) {
 		},
 	}
 
-	rfs := NewRebuiltForrest(nil, Graph{}, cbs)
+	t.Run("strict", func(t *testing.T) {
+		t.Parallel()
+		rfs := NewRebuiltForrest(nil, Graph{}, cbs, false)
 
-	tree, err := rfs.RebuiltTree(ctx, 305)
-	assert.EqualError(t, err, `failed to rebuild parent tree: 304: tree does not exist`)
-	assert.Nil(t, tree)
+		tree, err := rfs.RebuiltTree(ctx, 305)
+		assert.EqualError(t, err, `failed to rebuild parent tree: 304: tree does not exist`)
+		assert.Nil(t, tree)
+	})
+
+	t.Run("lax", func(t *testing.T) {
+		t.Parallel()
+		rfs := NewRebuiltForrest(nil, Graph{}, cbs, true)
+
+		tree, err := rfs.RebuiltTree(ctx, 305)
+		assert.NoError(t, err)
+		assert.NotNil(t, tree)
+	})
 }
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index e963f0a..3036938 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -27,6 +27,7 @@ type RebuiltTree struct {
 
 	rootErr      error
 	ancestorLoop bool
+	ancestorRoot btrfsprim.ObjID
 
 	ID        btrfsprim.ObjID
 	UUID      btrfsprim.UUID
@@ -127,6 +128,9 @@ func (tree *RebuiltTree) uncachedNodeIndex(ctx context.Context) rebuiltNodeIndex
 	}
 	for ancestor := tree; ancestor != nil; ancestor = ancestor.Parent {
 		indexer.idToTree[ancestor.ID] = ancestor
+		if ancestor.ID == tree.ancestorRoot {
+			break
+		}
 	}
 
 	ret := rebuiltNodeIndex{
@@ -261,11 +265,12 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 	// the "false"/failure case.  It will be called lots of times
 	// in a tight loop for both values that pass and values that
 	// fail.
+	root := tree.ancestorRoot
 	for {
 		if owner == tree.ID {
 			return true
 		}
-		if tree.Parent == nil || gen > tree.ParentGen {
+		if tree.Parent == nil || gen > tree.ParentGen || tree.ID == root {
 			return false
 		}
 		tree = tree.Parent
@@ -282,6 +287,7 @@ func (tree *RebuiltTree) isOwnerOK(owner btrfsprim.ObjID, gen btrfsprim.Generati
 //
 // When done with the map, call .RebuiltReleaseItems().
 func (tree *RebuiltTree) RebuiltAcquireItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+	tree.forrest.commitTrees(ctx, tree.ID)
 	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
@@ -303,6 +309,7 @@ func (tree *RebuiltTree) RebuiltReleaseItems() {
 //
 // When done with the map, call .RebuiltReleasePotentialItems().
 func (tree *RebuiltTree) RebuiltAcquirePotentialItems(ctx context.Context) *containers.SortedMap[btrfsprim.Key, ItemPtr] {
+	tree.forrest.commitTrees(ctx, tree.ID)
 	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
@@ -427,6 +434,15 @@ func (s rebuiltRootStats) String() string {
 // RebuiltAddRoot adds an additional root node to the tree.  It is
 // useful to call .RebuiltAddRoot() to re-attach part of the tree that
 // has been broken off.
+//
+// If the RebuiltForrest has laxAncestors=false, then:
+//
+//   - calls to RebuiltForrestExtendedCallbacks.AddedItem() are
+//     inhibited.
+//
+//   - calling RebuiltAddRoot on the ROOT_TREE or the UUID_TREE will
+//     panic if a tree other than the ROOT_TREE or UUID_TREE has been
+//     read from.
 func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.LogicalAddr) {
 	tree.mu.Lock()
 	defer tree.mu.Unlock()
@@ -436,7 +452,16 @@ func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.L
 
 	shouldFlush := tree.ID == btrfsprim.ROOT_TREE_OBJECTID || tree.ID == btrfsprim.UUID_TREE_OBJECTID
 
-	if extCB, ok := tree.forrest.cb.(RebuiltForrestExtendedCallbacks); ok {
+	if tree.forrest.laxAncestors && shouldFlush {
+		_ = tree.forrest.treesMu.Lock(ctx)
+		if tree.forrest.treesCommitted {
+			panic(fmt.Errorf("RebuiltTree(%v).RebuiltAddRoot called after a non-ROOT, non-UUID tree (%v) has been read from",
+				tree.ID, tree.forrest.treesCommitter))
+		}
+		tree.forrest.treesMu.Unlock()
+	}
+
+	if extCB, ok := tree.forrest.cb.(RebuiltForrestExtendedCallbacks); ok && !tree.forrest.laxAncestors {
 		var stats rebuiltRootStats
 		nodeToRoots := tree.acquireNodeIndex(ctx).nodeToRoots
 		stats.Nodes.D = len(nodeToRoots)
@@ -501,6 +526,7 @@ func (tree *RebuiltTree) RebuiltLeafToRoots(ctx context.Context, leaf btrfsvol.L
 			tree.ID, leaf))
 	}
 
+	tree.forrest.commitTrees(ctx, tree.ID)
 	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()
@@ -601,6 +627,7 @@ func (tree *RebuiltTree) TreeSubrange(ctx context.Context,
 
 // TreeWalk implements btrfstree.Tree.
 func (tree *RebuiltTree) TreeWalk(ctx context.Context, cbs btrfstree.TreeWalkHandler) {
+	tree.forrest.commitTrees(ctx, tree.ID)
 	tree.initRoots(ctx)
 	tree.mu.RLock()
 	defer tree.mu.RUnlock()