rebuildnodes: Support graceful shutdown

6 files changed, 81 insertions, 40 deletions

diff --git a/.golangci.yml b/.golangci.yml
index 6b9a830..355664f 100644
--- a/.golangci.yml
+++ b/.golangci.yml
@@ -51,7 +51,6 @@ linters:
     - revive
     - testpackage
     - thelper
-    - unparam
     - varnamelen
     - wrapcheck
 linters-settings:
diff --git a/cmd/btrfs-rec/inspect_rebuildnodes.go b/cmd/btrfs-rec/inspect_rebuildnodes.go
index e61e6d2..9c86c3a 100644
--- a/cmd/btrfs-rec/inspect_rebuildnodes.go
+++ b/cmd/btrfs-rec/inspect_rebuildnodes.go
@@ -33,21 +33,31 @@ func init() {
 			}
 			dlog.Infof(ctx, "... done reading %q", args[0])
 
-			rebuiltNodes, err := rebuildnodes.RebuildNodes(ctx, fs, nodeScanResults)
+			rebuilder, err := rebuildnodes.NewRebuilder(ctx, fs, nodeScanResults)
 			if err != nil {
 				return err
 			}
 
-			dlog.Info(ctx, "Writing re-built nodes to stdout...")
-			if err := writeJSONFile(os.Stdout, rebuiltNodes, lowmemjson.ReEncoder{
+			dlog.Info(ctx, "Rebuilding node tree...")
+			rebuildErr := rebuilder.Rebuild(ctx)
+			dst := os.Stdout
+			if rebuildErr != nil {
+				dst = os.Stderr
+				dlog.Errorf(ctx, "rebuild error: %v", rebuildErr)
+			}
+			dlog.Infof(ctx, "Writing re-built nodes to %s...", dst.Name())
+			if err := writeJSONFile(dst, rebuilder.ListRoots(), lowmemjson.ReEncoder{
 				Indent:                "\t",
 				ForceTrailingNewlines: true,
 			}); err != nil {
+				if rebuildErr != nil {
+					return rebuildErr
+				}
 				return err
 			}
 			dlog.Info(ctx, "... done writing")
 
-			return nil
+			return rebuildErr
 		},
 	})
 }
diff --git a/cmd/btrfs-rec/util.go b/cmd/btrfs-rec/util.go
index ffc03cc..9a0d60c 100644
--- a/cmd/btrfs-rec/util.go
+++ b/cmd/btrfs-rec/util.go
@@ -18,6 +18,7 @@ import (
 )
 
 type runeScanner struct {
+	ctx            context.Context //nolint:containedctx // For detecting shutdown from methods
 	progress       textui.Portion[int64]
 	progressWriter *textui.Progress[textui.Portion[int64]]
 	unreadCnt      uint64
@@ -31,6 +32,7 @@ func newRuneScanner(ctx context.Context, fh *os.File) (*runeScanner, error) {
 		return nil, err
 	}
 	ret := &runeScanner{
+		ctx: ctx,
 		progress: textui.Portion[int64]{
 			D: fi.Size(),
 		},
@@ -42,6 +44,9 @@ func newRuneScanner(ctx context.Context, fh *os.File) (*runeScanner, error) {
 }
 
 func (rs *runeScanner) ReadRune() (r rune, size int, err error) {
+	if err := rs.ctx.Err(); err != nil {
+		return 0, 0, err
+	}
 	r, size, err = rs.reader.ReadRune()
 	if rs.unreadCnt > 0 {
 		rs.unreadCnt--
@@ -53,8 +58,14 @@ func (rs *runeScanner) ReadRune() (r rune, size int, err error) {
 }
 
 func (rs *runeScanner) UnreadRune() error {
+	if err := rs.ctx.Err(); err != nil {
+		return err
+	}
+	if err := rs.reader.UnreadRune(); err != nil {
+		return err
+	}
 	rs.unreadCnt++
-	return rs.reader.UnreadRune()
+	return nil
 }
 
 func (rs *runeScanner) Close() error {
@@ -69,6 +80,9 @@ func readJSONFile[T any](ctx context.Context, filename string) (T, error) {
 		return zero, err
 	}
 	buf, err := newRuneScanner(dlog.WithField(ctx, "btrfs.read-json-file", filename), fh)
+	defer func() {
+		_ = buf.Close()
+	}()
 	if err != nil {
 		var zero T
 		return zero, err
@@ -78,7 +92,6 @@ func readJSONFile[T any](ctx context.Context, filename string) (T, error) {
 		var zero T
 		return zero, err
 	}
-	_ = buf.Close()
 	return ret, nil
 }
 
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
 }