Merge branch 'lukeshu/rebuilt-v2-pt1-naive'

5 files changed, 42 insertions, 26 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
index d1ccfac..7e2e49f 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
@@ -159,7 +159,7 @@ func (o *rebuilder) processTreeQueue(ctx context.Context) error {
 		}
 		// This will call o.AddedItem as nescessary, which
 		// inserts to o.addedItemQueue.
-		_ = o.rebuilt.RebuiltTree(ctx, o.curKey.TreeID)
+		_, _ = o.rebuilt.RebuiltTree(ctx, o.curKey.TreeID)
 	}
 
 	return nil
@@ -197,7 +197,7 @@ func (o *rebuilder) processAddedItemQueue(ctx context.Context) error {
 
 	for _, key := range queue {
 		ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.settle.item", key)
-		tree := o.rebuilt.RebuiltTree(ctx, key.TreeID)
+		tree := discardErr(o.rebuilt.RebuiltTree(ctx, key.TreeID))
 		incPtr, ok := tree.RebuiltAcquireItems(ctx).Load(key.Key)
 		tree.RebuiltReleaseItems()
 		if !ok {
@@ -333,9 +333,8 @@ func (o *rebuilder) sortSettledItemQueue(ctx context.Context, unorderedQueue con
 	// EXTENT_TREE; if that fails, then there can't be any items
 	// in the EXTENT_TREE for us to have to handle special, and
 	// all of the following code will fall through common-path.
-	extentTree := o.rebuilt.RebuiltTree(ctx, btrfsprim.EXTENT_TREE_OBJECTID)
 	var extentItems *containers.SortedMap[btrfsprim.Key, btrfsutil.ItemPtr]
-	if extentTree != nil {
+	if extentTree, err := o.rebuilt.RebuiltTree(ctx, btrfsprim.EXTENT_TREE_OBJECTID); err == nil {
 		extentItems = extentTree.RebuiltAcquireItems(ctx)
 		defer extentTree.RebuiltReleaseItems()
 	}
@@ -416,7 +415,7 @@ func (o *rebuilder) processSettledItemQueue(ctx context.Context) error {
 			ctx := dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.rebuild.process.item", key)
 			item := keyAndBody{
 				itemToVisit: key,
-				Body:        o.rebuilt.RebuiltTree(ctx, key.TreeID).ReadItem(ctx, key.Key),
+				Body:        discardErr(o.rebuilt.RebuiltTree(ctx, key.TreeID)).ReadItem(ctx, key.Key),
 			}
 			if key.TreeID == btrfsprim.EXTENT_TREE_OBJECTID &&
 				(key.ItemType == btrfsprim.EXTENT_ITEM_KEY || key.ItemType == btrfsprim.METADATA_ITEM_KEY) {
@@ -503,7 +502,7 @@ func (o *rebuilder) processAugmentQueue(ctx context.Context) error {
 			}
 			// This will call o.AddedItem as nescessary, which
 			// inserts to o.addedItemQueue.
-			o.rebuilt.RebuiltTree(ctx, treeID).RebuiltAddRoot(ctx, nodeAddr)
+			discardErr(o.rebuilt.RebuiltTree(ctx, treeID)).RebuiltAddRoot(ctx, nodeAddr)
 			progress.N++
 			progressWriter.Set(progress)
 		}
@@ -550,7 +549,7 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 		} else {
 			choices[choice] = ChoiceInfo{
 				Count:      1,
-				Distance:   discardOK(o.rebuilt.RebuiltTree(ctx, treeID).RebuiltCOWDistance(o.scan.Graph.Nodes[choice].Owner)),
+				Distance:   discardOK(discardErr(o.rebuilt.RebuiltTree(ctx, treeID)).RebuiltCOWDistance(o.scan.Graph.Nodes[choice].Owner)),
 				Generation: o.scan.Graph.Nodes[choice].Generation,
 			}
 		}
@@ -564,7 +563,7 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 			} else {
 				choices[choice] = ChoiceInfo{
 					Count:      1,
-					Distance:   discardOK(o.rebuilt.RebuiltTree(ctx, treeID).RebuiltCOWDistance(o.scan.Graph.Nodes[choice].Owner)),
+					Distance:   discardOK(discardErr(o.rebuilt.RebuiltTree(ctx, treeID)).RebuiltCOWDistance(o.scan.Graph.Nodes[choice].Owner)),
 					Generation: o.scan.Graph.Nodes[choice].Generation,
 				}
 			}
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go
index 0c52556..e227cc7 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_treecb.go
@@ -51,7 +51,7 @@ func (o forrestCallbacks) LookupRoot(ctx context.Context, tree btrfsprim.ObjID)
 		o.enqueueRetry(btrfsprim.ROOT_TREE_OBJECTID)
 		return 0, btrfsitem.Root{}, false
 	}
-	itemBody := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).ReadItem(ctx, foundKey)
+	itemBody := discardErr(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)).ReadItem(ctx, foundKey)
 	defer itemBody.Free()
 	switch itemBody := itemBody.(type) {
 	case *btrfsitem.Root:
@@ -77,7 +77,7 @@ func (o forrestCallbacks) LookupUUID(ctx context.Context, uuid btrfsprim.UUID) (
 		o.enqueueRetry(btrfsprim.UUID_TREE_OBJECTID)
 		return 0, false
 	}
-	itemBody := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID).ReadItem(ctx, wantKey.Key.Key())
+	itemBody := discardErr(o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)).ReadItem(ctx, wantKey.Key.Key())
 	defer itemBody.Free()
 	switch itemBody := itemBody.(type) {
 	case *btrfsitem.UUIDMap:
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
index 0526e93..382d88b 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild_wantcb.go
@@ -43,8 +43,8 @@ func (o graphCallbacks) Want(ctx context.Context, reason string, treeID btrfspri
 }
 
 func (o *rebuilder) _want(ctx context.Context, wantKey wantWithTree) (key btrfsprim.Key, ok bool) {
-	tree := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)
-	if tree == nil {
+	tree, err := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)
+	if err != nil {
 		o.enqueueRetry(wantKey.TreeID)
 		return btrfsprim.Key{}, false
 	}
@@ -98,8 +98,8 @@ func (o graphCallbacks) WantOff(ctx context.Context, reason string, treeID btrfs
 }
 
 func (o *rebuilder) _wantOff(ctx context.Context, wantKey wantWithTree) (ok bool) {
-	tree := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)
-	if tree == nil {
+	tree, err := o.rebuilt.RebuiltTree(ctx, wantKey.TreeID)
+	if err != nil {
 		o.enqueueRetry(wantKey.TreeID)
 		return false
 	}
@@ -144,8 +144,8 @@ func (o graphCallbacks) WantDirIndex(ctx context.Context, reason string, treeID
 	}
 	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
 
-	tree := o.rebuilt.RebuiltTree(ctx, treeID)
-	if tree == nil {
+	tree, err := o.rebuilt.RebuiltTree(ctx, treeID)
+	if err != nil {
 		o.enqueueRetry(treeID)
 		return
 	}
@@ -273,8 +273,8 @@ func (o graphCallbacks) _wantRange(
 	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
 	wantKey.Key.OffsetType = offsetRange
 
-	tree := o.rebuilt.RebuiltTree(ctx, treeID)
-	if tree == nil {
+	tree, err := o.rebuilt.RebuiltTree(ctx, treeID)
+	if err != nil {
 		o.enqueueRetry(treeID)
 		return
 	}
@@ -389,7 +389,7 @@ func (o graphCallbacks) WantCSum(ctx context.Context, reason string, inodeTree,
 		o.enqueueRetry(inodeTree)
 		return
 	}
-	tree := o.rebuilt.RebuiltTree(inodeCtx, inodeTree)
+	tree := discardErr(o.rebuilt.RebuiltTree(inodeCtx, inodeTree))
 	inodePtr, ok := tree.RebuiltAcquireItems(inodeCtx).Load(inodeWant.Key.Key())
 	tree.RebuiltReleaseItems()
 	if !ok {
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/util.go b/cmd/btrfs-rec/inspect/rebuildtrees/util.go
index 842fb55..ee06394 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/util.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/util.go
@@ -19,3 +19,7 @@ func roundUp[T constraints.Integer](n, d T) T {
 func discardOK[T any](val T, _ bool) T {
 	return val
 }
+
+func discardErr[T any](val T, _ error) T {
+	return val
+}
diff --git a/cmd/btrfs-rec/main.go b/cmd/btrfs-rec/main.go
index da3aced..26857f0 100644
--- a/cmd/btrfs-rec/main.go
+++ b/cmd/btrfs-rec/main.go
@@ -49,9 +49,10 @@ var globalFlags struct {
 	logLevel textui.LogLevelFlag
 	pvs      []string
 
-	mappings string
-	nodeList string
-	rebuild  bool
+	mappings  string
+	nodeList  string
+	rebuildV1 bool
+	rebuildV2 bool
 
 	stopProfiling profile.StopFunc
 
@@ -101,8 +102,11 @@ func main() {
 		"load node list (output of 'btrfs-recs inspect [rebuild-mappings] list-nodes') from external JSON file `nodes.json`")
 	noError(argparser.MarkPersistentFlagFilename("node-list"))
 
-	argparser.PersistentFlags().BoolVar(&globalFlags.rebuild, "rebuild", false,
-		"attempt to rebuild broken btrees when reading")
+	argparser.PersistentFlags().BoolVar(&globalFlags.rebuildV1, "rebuild", false,
+		"attempt to rebuild broken btrees when reading (v1)")
+
+	argparser.PersistentFlags().BoolVar(&globalFlags.rebuildV2, "rebuild-v2", false,
+		"attempt to rebuild broken btrees when reading (v2)")
 
 	globalFlags.stopProfiling = profile.AddProfileFlags(argparser.PersistentFlags(), "profile.")
 
@@ -210,15 +214,24 @@ func runWithRawFSAndNodeList(runE func(*btrfs.FS, []btrfsvol.LogicalAddr, *cobra
 func _runWithReadableFS(wantNodeList bool, runE func(btrfs.ReadableFS, []btrfsvol.LogicalAddr, *cobra.Command, []string) error) func(*cobra.Command, []string) error {
 	inner := func(fs *btrfs.FS, nodeList []btrfsvol.LogicalAddr, cmd *cobra.Command, args []string) error {
 		var rfs btrfs.ReadableFS = fs
-		if globalFlags.rebuild {
+		if globalFlags.rebuildV1 {
 			rfs = btrfsutil.NewOldRebuiltForrest(fs)
+		} else if globalFlags.rebuildV2 {
+			ctx := cmd.Context()
+
+			graph, err := btrfsutil.ReadGraph(ctx, fs, nodeList)
+			if err != nil {
+				return err
+			}
+
+			rfs = btrfsutil.NewRebuiltForrest(fs, graph, nil)
 		}
 
 		return runE(rfs, nodeList, cmd, args)
 	}
 
 	return func(cmd *cobra.Command, args []string) error {
-		if wantNodeList {
+		if wantNodeList || globalFlags.rebuildV2 {
 			return runWithRawFSAndNodeList(inner)(cmd, args)
 		}
 		return runWithRawFS(func(fs *btrfs.FS, cmd *cobra.Command, args []string) error {