Merge branch 'lukeshu/rebuild-nodes-take4'

2 files changed, 46 insertions, 12 deletions

diff --git a/cmd/btrfs-rec/inspect_rebuildnodes.go b/cmd/btrfs-rec/inspect_rebuildnodes.go
index e61e6d2..d813f36 100644
--- a/cmd/btrfs-rec/inspect_rebuildnodes.go
+++ b/cmd/btrfs-rec/inspect_rebuildnodes.go
@@ -5,7 +5,10 @@
 package main
 
 import (
+	"context"
 	"os"
+	"runtime"
+	"time"
 
 	"git.lukeshu.com/go/lowmemjson"
 	"github.com/datawire/dlib/dlog"
@@ -15,6 +18,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 func init() {
@@ -26,28 +30,45 @@ func init() {
 		RunE: func(fs *btrfs.FS, cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 
-			dlog.Infof(ctx, "Reading %q...", args[0])
-			nodeScanResults, err := readJSONFile[btrfsinspect.ScanDevicesResult](ctx, args[0])
-			if err != nil {
-				return err
-			}
-			dlog.Infof(ctx, "... done reading %q", args[0])
+			// This is wrapped in a func in order to *ensure* that `nodeScanResults` goes out of scope once
+			// `rebuilder` has been created.
+			rebuilder, err := func(ctx context.Context) (rebuildnodes.Rebuilder, error) {
+				dlog.Infof(ctx, "Reading %q...", args[0])
+				nodeScanResults, err := readJSONFile[btrfsinspect.ScanDevicesResult](ctx, args[0])
+				if err != nil {
+					return nil, err
+				}
+				dlog.Infof(ctx, "... done reading %q", args[0])
 
-			rebuiltNodes, err := rebuildnodes.RebuildNodes(ctx, fs, nodeScanResults)
+				return rebuildnodes.NewRebuilder(ctx, fs, nodeScanResults)
+			}(ctx)
 			if err != nil {
 				return err
 			}
 
-			dlog.Info(ctx, "Writing re-built nodes to stdout...")
-			if err := writeJSONFile(os.Stdout, rebuiltNodes, lowmemjson.ReEncoder{
+			runtime.GC()
+			time.Sleep(textui.LiveMemUseUpdateInterval) // let the logs reflect that GC right away
+
+			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
 }