Expose node-lists as a thing on the CLI

7 files changed, 214 insertions, 47 deletions

diff --git a/cmd/btrfs-rec/inspect_listnodes.go b/cmd/btrfs-rec/inspect_listnodes.go
new file mode 100644
index 0000000..d9b24ed
--- /dev/null
+++ b/cmd/btrfs-rec/inspect_listnodes.go
@@ -0,0 +1,51 @@
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package main
+
+import (
+	"os"
+
+	"git.lukeshu.com/go/lowmemjson"
+	"github.com/datawire/dlib/dlog"
+	"github.com/datawire/ocibuild/pkg/cliutil"
+	"github.com/spf13/cobra"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsutil"
+)
+
+func init() {
+	inspectors.AddCommand(&cobra.Command{
+		Use:   "list-nodes",
+		Short: "Scan the filesystem for btree nodes",
+		Long: "" +
+			"This scans the filesystem sector-by-sector looking for nodes.  " +
+			"If you are needing to rebuild the chunk/dev-extent/blockgroup " +
+			"trees with `btrfs-rec inspect rebuild-mappings` anyway, you may " +
+			"want to instead use `btrfs-rec inspect rebuild-mappings list-nodes` " +
+			"to take advantage of the sector-by-sector scan that's already " +
+			"performed by `btrfs-rec inspect rebuild-mappings scan`.",
+		Args: cliutil.WrapPositionalArgs(cobra.ExactArgs(1)),
+		RunE: runWithRawFS(func(fs *btrfs.FS, cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			nodeList, err := btrfsutil.ListNodes(ctx, fs)
+			if err != nil {
+				return err
+			}
+
+			dlog.Infof(ctx, "Writing nodes to stdout...")
+			if err := writeJSONFile(os.Stdout, nodeList, lowmemjson.ReEncoderConfig{
+				Indent:                "\t",
+				ForceTrailingNewlines: true,
+			}); err != nil {
+				return err
+			}
+			dlog.Info(ctx, "... done writing")
+
+			return nil
+		}),
+	})
+}
diff --git a/cmd/btrfs-rec/inspect_lstrees.go b/cmd/btrfs-rec/inspect_lstrees.go
index 1ff7671..05c3a57 100644
--- a/cmd/btrfs-rec/inspect_lstrees.go
+++ b/cmd/btrfs-rec/inspect_lstrees.go
@@ -26,14 +26,24 @@ import (
 )
 
 func init() {
-	var scandevicesFilename string
+	var nodeListFilename string
 	cmd := &cobra.Command{
 		Use:   "ls-trees",
 		Short: "A brief view what types of items are in each tree",
-		Args:  cliutil.WrapPositionalArgs(cobra.NoArgs),
+		Long: "" +
+			"If no --node-list is given, then a slow sector-by-sector scan " +
+			"will be used to find all lost+found nodes.",
+		Args: cliutil.WrapPositionalArgs(cobra.NoArgs),
 		RunE: runWithRawFS(func(fs *btrfs.FS, cmd *cobra.Command, _ []string) error {
 			ctx := cmd.Context()
-			nodeList, err := readNodeList(ctx, scandevicesFilename)
+
+			var nodeList []btrfsvol.LogicalAddr
+			var err error
+			if nodeListFilename != "" {
+				nodeList, err = readJSONFile[[]btrfsvol.LogicalAddr](ctx, nodeListFilename)
+			} else {
+				nodeList, err = btrfsutil.ListNodes(ctx, fs)
+			}
 			if err != nil {
 				return err
 			}
@@ -86,7 +96,7 @@ func init() {
 				},
 			})
 
-			if scandevicesFilename != "" {
+			{
 				treeErrCnt = 0
 				treeItemCnt = make(map[btrfsitem.Type]int)
 				textui.Fprintf(os.Stdout, "lost+found\n")
@@ -114,8 +124,9 @@ func init() {
 			return nil
 		}),
 	}
-	cmd.Flags().StringVar(&scandevicesFilename, "scandevices", "", "Output of 'btrfs-recs inspect rebuild-mappings scan' to use for a lost+found tree")
-	noError(cmd.MarkFlagFilename("scandevices"))
+	cmd.Flags().StringVar(&nodeListFilename, "node-list", "",
+		"Output of 'btrfs-recs inspect [rebuild-mappings] list-nodes' to use for a lost+found tree")
+	noError(cmd.MarkFlagFilename("node-list"))
 
 	inspectors.AddCommand(cmd)
 }
diff --git a/cmd/btrfs-rec/inspect_rebuildmappings.go b/cmd/btrfs-rec/inspect_rebuildmappings.go
index 9c6b1cc..b215e7a 100644
--- a/cmd/btrfs-rec/inspect_rebuildmappings.go
+++ b/cmd/btrfs-rec/inspect_rebuildmappings.go
@@ -14,6 +14,9 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/cmd/btrfs-rec/inspect/rebuildmappings"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
 func init() {
@@ -118,5 +121,46 @@ func init() {
 		}),
 	})
 
+	cmd.AddCommand(&cobra.Command{
+		Use:   "list-nodes",
+		Short: "Produce a listing of btree nodes from previously read data",
+		Long: "" +
+			"This is a variant of `btrfs-rec inspect list-nodes` that takes " +
+			"advantage of using previously read data from " +
+			"`btrfs-rec inspect rebuild-nodes scan`.",
+		Args: cliutil.WrapPositionalArgs(cobra.ExactArgs(1)),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			scanResults, err := readJSONFile[rebuildmappings.ScanDevicesResult](ctx, args[0])
+			if err != nil {
+				return err
+			}
+
+			var cnt int
+			for _, devResults := range scanResults {
+				cnt += len(devResults.FoundNodes)
+			}
+			set := make(containers.Set[btrfsvol.LogicalAddr], cnt)
+			for _, devResults := range scanResults {
+				for laddr := range devResults.FoundNodes {
+					set.Insert(laddr)
+				}
+			}
+			nodeList := maps.SortedKeys(set)
+
+			dlog.Infof(ctx, "Writing nodes to stdout...")
+			if err := writeJSONFile(os.Stdout, nodeList, lowmemjson.ReEncoderConfig{
+				Indent:                "\t",
+				ForceTrailingNewlines: true,
+			}); err != nil {
+				return err
+			}
+			dlog.Info(ctx, "... done writing")
+
+			return nil
+		},
+	})
+
 	inspectors.AddCommand(cmd)
 }
diff --git a/cmd/btrfs-rec/inspect_rebuildtrees.go b/cmd/btrfs-rec/inspect_rebuildtrees.go
index f66f287..676533a 100644
--- a/cmd/btrfs-rec/inspect_rebuildtrees.go
+++ b/cmd/btrfs-rec/inspect_rebuildtrees.go
@@ -5,7 +5,6 @@
 package main
 
 import (
-	"context"
 	"os"
 	"runtime"
 	"time"
@@ -17,25 +16,39 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/cmd/btrfs-rec/inspect/rebuildtrees"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
 func init() {
-	inspectors.AddCommand(&cobra.Command{
-		Use:  "rebuild-trees NODESCAN.json",
-		Args: cliutil.WrapPositionalArgs(cobra.ExactArgs(1)),
+	var nodeListFilename string
+	cmd := &cobra.Command{
+		Use: "rebuild-trees",
+		Long: "" +
+			"Rebuild broken btrees based on missing items that are implied " +
+			"by present items.  This requires functioning " +
+			"chunk/dev-extent/blockgroup trees, which can be rebuilt " +
+			"separately with `btrfs-rec inspect rebuild-mappings`.\n" +
+			"\n" +
+			"If no --node-list is given, then a slow sector-by-sector scan " +
+			"will be used to find all nodes.",
+		Args: cliutil.WrapPositionalArgs(cobra.NoArgs),
 		RunE: runWithRawFS(func(fs *btrfs.FS, cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 
-			// This is wrapped in a func in order to *ensure* that `nodeList` goes out of scope once
-			// `rebuilder` has been created.
-			rebuilder, err := func(ctx context.Context) (rebuildtrees.Rebuilder, error) {
-				nodeList, err := readNodeList(ctx, args[0])
-				if err != nil {
-					return nil, err
-				}
-				return rebuildtrees.NewRebuilder(ctx, fs, nodeList)
-			}(ctx)
+			var nodeList []btrfsvol.LogicalAddr
+			var err error
+			if nodeListFilename != "" {
+				nodeList, err = readJSONFile[[]btrfsvol.LogicalAddr](ctx, nodeListFilename)
+			} else {
+				nodeList, err = btrfsutil.ListNodes(ctx, fs)
+			}
+			if err != nil {
+				return err
+			}
+
+			rebuilder, err := rebuildtrees.NewRebuilder(ctx, fs, nodeList)
 			if err != nil {
 				return err
 			}
@@ -65,5 +78,10 @@ func init() {
 
 			return rebuildErr
 		}),
-	})
+	}
+	cmd.Flags().StringVar(&nodeListFilename, "node-list", "",
+		"Output of 'btrfs-recs inspect [rebuild-mappings] list-nodes' to use for the node list")
+	noError(cmd.MarkFlagFilename("node-list"))
+
+	inspectors.AddCommand(cmd)
 }
diff --git a/cmd/btrfs-rec/util.go b/cmd/btrfs-rec/util.go
index 2e4bb84..3d751a6 100644
--- a/cmd/btrfs-rec/util.go
+++ b/cmd/btrfs-rec/util.go
@@ -13,10 +13,6 @@ import (
 	"git.lukeshu.com/go/lowmemjson"
 	"github.com/datawire/dlib/dlog"
 
-	"git.lukeshu.com/btrfs-progs-ng/cmd/btrfs-rec/inspect/rebuildmappings"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 	"git.lukeshu.com/btrfs-progs-ng/lib/streamio"
 )
 
@@ -51,24 +47,3 @@ func writeJSONFile(w io.Writer, obj any, cfg lowmemjson.ReEncoderConfig) (err er
 	}()
 	return lowmemjson.NewEncoder(lowmemjson.NewReEncoder(buffer, cfg)).Encode(obj)
 }
-
-func readNodeList(ctx context.Context, filename string) ([]btrfsvol.LogicalAddr, error) {
-	if filename == "" {
-		return nil, nil
-	}
-	results, err := readJSONFile[rebuildmappings.ScanDevicesResult](ctx, filename)
-	if err != nil {
-		return nil, err
-	}
-	var cnt int
-	for _, devResults := range results {
-		cnt += len(devResults.FoundNodes)
-	}
-	set := make(containers.Set[btrfsvol.LogicalAddr], cnt)
-	for _, devResults := range results {
-		for laddr := range devResults.FoundNodes {
-			set.Insert(laddr)
-		}
-	}
-	return maps.SortedKeys(set), nil
-}
diff --git a/lib/btrfsutil/listnodes.go b/lib/btrfsutil/listnodes.go
new file mode 100644
index 0000000..16300da
--- /dev/null
+++ b/lib/btrfsutil/listnodes.go
@@ -0,0 +1,66 @@
+// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+import (
+	"context"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+type nodeScanner struct {
+	nodes containers.Set[btrfsvol.LogicalAddr]
+}
+
+type nodeStats struct {
+	numNodes int
+}
+
+func (s nodeStats) String() string {
+	return textui.Sprintf("found: %d nodes", s.numNodes)
+}
+
+var _ DeviceScanner[nodeStats, containers.Set[btrfsvol.LogicalAddr]] = (*nodeScanner)(nil)
+
+func newNodeScanner(ctx context.Context, sb btrfstree.Superblock, numBytes btrfsvol.PhysicalAddr, numSectors int) DeviceScanner[nodeStats, containers.Set[btrfsvol.LogicalAddr]] {
+	s := new(nodeScanner)
+	s.nodes = make(containers.Set[btrfsvol.LogicalAddr])
+	return s
+}
+
+func (s *nodeScanner) ScanStats() nodeStats {
+	return nodeStats{numNodes: len(s.nodes)}
+}
+
+func (*nodeScanner) ScanSector(ctx context.Context, dev *btrfs.Device, paddr btrfsvol.PhysicalAddr) error {
+	return nil
+}
+
+func (s *nodeScanner) ScanNode(ctx context.Context, nodeRef *diskio.Ref[btrfsvol.PhysicalAddr, btrfstree.Node]) error {
+	s.nodes.Insert(nodeRef.Data.Head.Addr)
+	return nil
+}
+
+func (s *nodeScanner) ScanDone(ctx context.Context) (containers.Set[btrfsvol.LogicalAddr], error) {
+	return s.nodes, nil
+}
+
+func ListNodes(ctx context.Context, fs *btrfs.FS) ([]btrfsvol.LogicalAddr, error) {
+	perDev, err := ScanDevices[nodeStats, containers.Set[btrfsvol.LogicalAddr]](ctx, fs, newNodeScanner)
+	if err != nil {
+		return nil, err
+	}
+	set := make(containers.Set[btrfsvol.LogicalAddr])
+	for _, devSet := range perDev {
+		set.InsertFrom(devSet)
+	}
+	return maps.SortedKeys(set), nil
+}
diff --git a/scripts/main.sh b/scripts/main.sh
index 2ea5213..dddb9bb 100755
--- a/scripts/main.sh
+++ b/scripts/main.sh
@@ -33,6 +33,8 @@ export GOMEMLIMIT="$(awk '/^MemTotal:/{ print $2 "KiB" }' </proc/meminfo)"
 
 run-btrfs-rec $gendir/0.scandevices.json \
     inspect rebuild-mappings scan
+run-btrfs-rec $gendir/0.nodes.json \
+    inspect rebuild-mappings list-nodes $gendir/0.scandevices.json
 run-btrfs-rec $gendir/1.mappings.json \
     inspect rebuild-mappings process $gendir/0.scandevices.json
 
@@ -62,11 +64,11 @@ run-btrfs-rec $gendir/2.mappings.json \
 
 run-btrfs-rec $gendir/3.trees.json \
     --mappings=$gendir/2.mappings.json \
-    inspect rebuild-trees $gendir/0.scandevices.json
+    inspect rebuild-trees --node-list=$gendir/0.nodes.json
 
 run-btrfs-rec $gendir/4.ls-files.txt \
     --mappings=$gendir/2.mappings.json \
     inspect ls-files
 run-btrfs-rec $gendir/4.ls-trees.txt \
     --mappings=$gendir/2.mappings.json \
-    inspect ls-trees --scandevices=$gendir/0.scandevices.json
+    inspect ls-trees --node-list=$gendir/0.nodes.json