Merge branch 'lukeshu/rebuild-nodes-take4'

1 files changed, 11 insertions, 7 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 7e96e29..7e19802 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -20,14 +20,15 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, *keyio.Handle, error) {
-	dlog.Infof(ctx, "Reading node data from FS...")
-
+func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (btrfstree.Superblock, graph.Graph, *keyio.Handle, error) {
+	dlog.Info(ctx, "Reading superblock...")
 	sb, err := fs.Superblock()
 	if err != nil {
-		return graph.Graph{}, nil, err
+		return btrfstree.Superblock{}, graph.Graph{}, nil, err
 	}
 
+	dlog.Infof(ctx, "Reading node data from FS...")
+
 	var stats textui.Portion[int]
 	stats.D = countNodes(scanResults)
 	progressWriter := textui.NewProgress[textui.Portion[int]](
@@ -40,11 +41,14 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	progressWriter.Set(stats)
 	for _, devResults := range scanResults {
 		for laddr := range devResults.FoundNodes {
+			if err := ctx.Err(); err != nil {
+				return btrfstree.Superblock{}, graph.Graph{}, nil, err
+			}
 			nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{
 				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
 			})
 			if err != nil {
-				return graph.Graph{}, nil, err
+				return btrfstree.Superblock{}, graph.Graph{}, nil, err
 			}
 
 			nodeGraph.InsertNode(nodeRef)
@@ -61,9 +65,9 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	dlog.Info(ctx, "... done reading node data")
 
 	if err := nodeGraph.FinalCheck(ctx, fs, *sb); err != nil {
-		return graph.Graph{}, nil, err
+		return btrfstree.Superblock{}, graph.Graph{}, nil, err
 	}
 	keyIO.SetGraph(*nodeGraph)
 
-	return *nodeGraph, keyIO, nil
+	return *sb, *nodeGraph, keyIO, nil
 }