From 8efc82d0b1a167830970135c78d173667080b116 Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Wed, 28 Dec 2022 18:49:09 -0700
Subject: rebuildnodes: Support graceful shutdown

---
 lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

(limited to 'lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go')

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
 }
-- 
cgit v1.2.3-2-g168b