rebuildnodes: Support graceful shutdown

3 files changed, 52 insertions, 33 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
index b1ae7be..ffe225a 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go
@@ -305,7 +305,7 @@ func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
 
 	progress()
 	for _, node := range maps.SortedKeys(graph.Nodes) {
-		tree.indexNode(graph, node, nodeToRoots, progress, nil)
+		tree.indexNode(ctx, graph, node, nodeToRoots, progress, nil)
 	}
 	progressWriter.Done()
 
@@ -317,8 +317,11 @@ func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) {
 	}
 }
 
-func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
+func (tree *rebuiltTree) indexNode(ctx context.Context, graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) {
 	defer progress()
+	if err := ctx.Err(); err != nil {
+		return
+	}
 	if _, done := index[node]; done {
 		return
 	}
@@ -339,7 +342,7 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd
 		return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner, graph.Nodes[kp.FromNode].Generation)
 	})
 	for _, kp := range kps {
-		tree.indexNode(graph, kp.FromNode, index, progress, stack)
+		tree.indexNode(ctx, graph, kp.FromNode, index, progress, stack)
 		if len(index[kp.FromNode]) > 0 {
 			if roots == nil {
 				roots = make(containers.Set[btrfsvol.LogicalAddr])
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index b6c1359..26f2a44 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -44,47 +44,38 @@ func (o keyAndTree) String() string {
 }
 
 type rebuilder struct {
-	sb      btrfstree.Superblock
-	rebuilt *btrees.RebuiltTrees
-
+	sb    btrfstree.Superblock
 	graph graph.Graph
 	keyIO *keyio.Handle
 
+	rebuilt *btrees.RebuiltTrees
+
 	curKey       keyAndTree
 	treeQueue    containers.Set[btrfsprim.ObjID]
 	itemQueue    containers.Set[keyAndTree]
 	augmentQueue map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int
 }
 
-func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) {
-	_ctx := ctx
+type Rebuilder interface {
+	Rebuild(context.Context) error
+	ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr]
+}
 
-	ctx = dlog.WithField(_ctx, "btrfsinspect.rebuild-nodes.step", "read-fs-data")
-	dlog.Info(ctx, "Reading superblock...")
-	sb, err := fs.Superblock()
-	if err != nil {
-		return nil, err
-	}
-	nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
+func NewRebuilder(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (Rebuilder, error) {
+	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.step", "read-fs-data")
+	sb, nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging
 	if err != nil {
 		return nil, err
 	}
 
-	ctx = dlog.WithField(_ctx, "btrfsinspect.rebuild-nodes.step", "rebuild")
-	dlog.Info(ctx, "Rebuilding node tree...")
 	o := &rebuilder{
-		sb: *sb,
-
+		sb:    sb,
 		graph: nodeGraph,
 		keyIO: keyIO,
 	}
-	o.rebuilt = btrees.NewRebuiltTrees(*sb, nodeGraph, keyIO,
+	o.rebuilt = btrees.NewRebuiltTrees(sb, nodeGraph, keyIO,
 		o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID)
-	if err := o.rebuild(ctx); err != nil {
-		return nil, err
-	}
-
-	return o.rebuilt.ListRoots(), nil
+	return o, nil
 }
 
 func (o *rebuilder) ioErr(ctx context.Context, err error) {
@@ -93,7 +84,13 @@ func (o *rebuilder) ioErr(ctx context.Context, err error) {
 	panic(err)
 }
 
-func (o *rebuilder) rebuild(_ctx context.Context) error {
+func (o *rebuilder) ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] {
+	return o.rebuilt.ListRoots()
+}
+
+func (o *rebuilder) Rebuild(_ctx context.Context) error {
+	_ctx = dlog.WithField(_ctx, "btrfsinspect.rebuild-nodes.step", "rebuild")
+
 	// Initialize
 	o.itemQueue = make(containers.Set[keyAndTree])
 	o.augmentQueue = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int)
@@ -116,6 +113,9 @@ func (o *rebuilder) rebuild(_ctx context.Context) error {
 		// Because trees can be wildly different sizes, it's impossible to have a meaningful
 		// progress percentage here.
 		for _, treeID := range maps.SortedKeys(treeQueue) {
+			if err := _ctx.Err(); err != nil {
+				return err
+			}
 			o.rebuilt.AddTree(stepCtx, treeID)
 		}
 
@@ -134,6 +134,10 @@ func (o *rebuilder) rebuild(_ctx context.Context) error {
 			itemCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.process.item", key)
 			progress.N = i
 			progressWriter.Set(progress)
+			if err := _ctx.Err(); err != nil {
+				progressWriter.Done()
+				return err
+			}
 			o.curKey = key
 			itemBody, ok := o.rebuilt.Load(itemCtx, key.TreeID, key.Key)
 			if !ok {
@@ -157,6 +161,9 @@ func (o *rebuilder) rebuild(_ctx context.Context) error {
 		progress.N = 0
 		progress.D = 0
 		for _, treeID := range maps.SortedKeys(o.augmentQueue) {
+			if err := _ctx.Err(); err != nil {
+				return err
+			}
 			treeCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.augment.tree", treeID)
 			resolvedAugments[treeID] = o.resolveTreeAugments(treeCtx, o.augmentQueue[treeID])
 			progress.D += len(resolvedAugments[treeID])
@@ -167,6 +174,11 @@ func (o *rebuilder) rebuild(_ctx context.Context) error {
 		for _, treeID := range maps.SortedKeys(resolvedAugments) {
 			treeCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.augment.tree", treeID)
 			for _, nodeAddr := range maps.SortedKeys(resolvedAugments[treeID]) {
+				if err := _ctx.Err(); err != nil {
+					progressWriter.Set(progress)
+					progressWriter.Done()
+					return err
+				}
 				progressWriter.Set(progress)
 				o.rebuilt.AddRoot(treeCtx, treeID, nodeAddr)
 				progress.N++
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
 }