wip

4 files changed, 72 insertions, 153 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
index 941ea4c..a9ff24d 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
@@ -31,17 +31,12 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 		uuidMap: uuidMap,
 	}
 
-	foundRoots, err := lostAndFoundNodes(ctx, nfs, nodeScanResults)
+	orphanedNodes, rebuiltNodes, err := classifyNodes(ctx, nfs, nodeScanResults)
 	if err != nil {
 		return nil, err
 	}
 
-	rebuiltNodes, err := reInitBrokenNodes(ctx, nfs, nodeScanResults, foundRoots)
-	if err != nil {
-		return nil, err
-	}
-
-	if err := reAttachNodes(ctx, nfs, foundRoots, rebuiltNodes); err != nil {
+	if err := reAttachNodes(ctx, nfs, orphanedNodes, rebuiltNodes); err != nil {
 		return nil, err
 	}
 
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reinit.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
index 49bc989..9fb5bf5 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reinit.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
@@ -6,9 +6,11 @@ package rebuildnodes
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	iofs "io/fs"
 	"reflect"
+	"strings"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -19,24 +21,34 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
-func reInitBrokenNodes(ctx context.Context, fs _FS, nodeScanResults btrfsinspect.ScanDevicesResult, foundRoots map[btrfsvol.LogicalAddr]struct{}) (map[btrfsvol.LogicalAddr]*RebuiltNode, error) {
-	dlog.Info(ctx, "Initializing nodes to re-build...")
+// classifyNodes returns
+//
+//  1. the set of nodes don't have another node claiming it as a child, and
+//  2. the set of bad-nodes, with reconstructed headers filled in.
+func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDevicesResult) (
+	orphanedNodes map[btrfsvol.LogicalAddr]struct{},
+	rebuiltNodes map[btrfsvol.LogicalAddr]*RebuiltNode,
+	err error,
+) {
+	dlog.Info(ctx, "Walking trees to identify orphan and broken nodes...")
 
 	sb, err := fs.Superblock()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-
 	chunkTreeUUID, ok := getChunkTreeUUID(ctx, fs)
 	if !ok {
-		return nil, fmt.Errorf("could not look up chunk tree UUID")
+		return nil, nil, fmt.Errorf("could not look up chunk tree UUID")
 	}
 
 	lastPct := -1
-	total := countNodes(nodeScanResults)
-	progress := func(done int) {
+	total := countNodes(scanResults)
+	visitedNodes := make(map[btrfsvol.LogicalAddr]struct{})
+	progress := func() {
+		done := len(visitedNodes)
 		pct := int(100 * float64(done) / float64(total))
 		if pct != lastPct || done == total {
 			dlog.Infof(ctx, "... %v%% (%v/%v)",
@@ -45,18 +57,21 @@ func reInitBrokenNodes(ctx context.Context, fs _FS, nodeScanResults btrfsinspect
 		}
 	}
 
-	rebuiltNodes := make(map[btrfsvol.LogicalAddr]*RebuiltNode)
-	visitedNodes := make(map[btrfsvol.LogicalAddr]struct{})
+	rebuiltNodes = make(map[btrfsvol.LogicalAddr]*RebuiltNode)
 	walkHandler := btrfstree.TreeWalkHandler{
-		Node: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
+		PreNode: func(path btrfstree.TreePath) error {
 			addr := path.Node(-1).ToNodeAddr
 			if _, alreadyVisited := visitedNodes[addr]; alreadyVisited {
 				// Can happen because of COW subvolumes;
 				// this is really a DAG not a tree.
 				return iofs.SkipDir
 			}
+			return nil
+		},
+		Node: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
+			addr := path.Node(-1).ToNodeAddr
 			visitedNodes[addr] = struct{}{}
-			progress(len(visitedNodes))
+			progress()
 			return nil
 		},
 		BadNode: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
@@ -98,29 +113,54 @@ func reInitBrokenNodes(ctx context.Context, fs _FS, nodeScanResults btrfsinspect
 		},
 	}
 
-	// We use WalkAllTrees instead of just iterating over
-	// nodeScanResults so that we don't need to specifically check
-	// if any of the root nodes referenced directly by the
-	// superblock are dead.
-	progress(len(visitedNodes))
 	btrfsutil.WalkAllTrees(ctx, fs, btrfsutil.WalkAllTreesHandler{
+		TreeWalkHandler: walkHandler,
 		Err: func(err *btrfsutil.WalkError) {
 			// do nothing
+			if !errors.Is(err, btrfstree.ErrNotANode) && !strings.Contains(err.Error(), "read: could not map logical address") {
+				dlog.Errorf(ctx, "dbg walk err: %v", err)
+			}
 		},
-		TreeWalkHandler: walkHandler,
 	})
-	for foundRoot := range foundRoots {
-		walkFromNode(ctx, fs, foundRoot,
+
+	// Start with 'orphanedRoots' as a complete set of all orphaned nodes, and then delete
+	// non-root entries from it.
+	orphanedNodes = make(map[btrfsvol.LogicalAddr]struct{})
+	for _, devResults := range scanResults {
+		for laddr := range devResults.FoundNodes {
+			if _, attached := visitedNodes[laddr]; !attached {
+				orphanedNodes[laddr] = struct{}{}
+			}
+		}
+	}
+	if len(visitedNodes)+len(orphanedNodes) != total {
+		panic("should not happen")
+	}
+	dlog.Infof(ctx,
+		"... (finished processing %v attached nodes, proceeding to process %v lost nodes, for a total of %v)",
+		len(visitedNodes), len(orphanedNodes), len(visitedNodes)+len(orphanedNodes))
+	for _, potentialRoot := range maps.SortedKeys(orphanedNodes) {
+		walkHandler.Node = func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
+			addr := path.Node(-1).ToNodeAddr
+			if addr != potentialRoot {
+				delete(orphanedNodes, addr)
+			}
+			visitedNodes[addr] = struct{}{}
+			progress()
+			return nil
+		}
+		walkFromNode(ctx, fs, potentialRoot,
 			func(err *btrfstree.TreeError) {
 				// do nothing
 			},
-			walkHandler)
+			walkHandler,
+		)
 	}
 
 	if len(visitedNodes) != total {
 		panic("should not happen")
 	}
 
-	dlog.Infof(ctx, "... initialized %d nodes", len(rebuiltNodes))
-	return rebuiltNodes, nil
+	dlog.Infof(ctx, "... identified %d orphaned nodes and re-built %d nodes", len(orphanedNodes), len(rebuiltNodes))
+	return orphanedNodes, rebuiltNodes, nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_lostandfound.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_lostandfound.go
deleted file mode 100644
index 4efaeab..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_lostandfound.go
+++ /dev/null
@@ -1,116 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package rebuildnodes
-
-import (
-	"context"
-	"errors"
-	iofs "io/fs"
-	"strings"
-
-	"github.com/datawire/dlib/dlog"
-
-	"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/btrfsprogs/btrfsinspect"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
-)
-
-// lostAndFoundNodes returns the set of nodes don't have another node
-// claiming it as a child.
-func lostAndFoundNodes(ctx context.Context, fs _FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsvol.LogicalAddr]struct{}, error) {
-	dlog.Info(ctx, "Identifying lost+found nodes...")
-
-	lastPct := -1
-	total := countNodes(nodeScanResults)
-	progress := func(done int) {
-		pct := int(100 * float64(done) / float64(total))
-		if pct != lastPct || done == total {
-			dlog.Infof(ctx, "... %v%% (%v/%v)",
-				pct, done, total)
-			lastPct = pct
-		}
-	}
-
-	visitedNodes := make(map[btrfsvol.LogicalAddr]struct{})
-	btrfsutil.WalkAllTrees(ctx, fs, btrfsutil.WalkAllTreesHandler{
-		TreeWalkHandler: btrfstree.TreeWalkHandler{
-			// Don't use `PreNode` because we don't want to run this on bad nodes.
-			Node: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
-				addr := path.Node(-1).ToNodeAddr
-				if _, alreadyVisited := visitedNodes[addr]; alreadyVisited {
-					// Can happen because of COW subvolumes;
-					// this is really a DAG not a tree.
-					return iofs.SkipDir
-				}
-				visitedNodes[addr] = struct{}{}
-				progress(len(visitedNodes))
-				return nil
-			},
-		},
-		Err: func(err *btrfsutil.WalkError) {
-			// do nothing
-			if !errors.Is(err, btrfstree.ErrNotANode) && !strings.Contains(err.Error(), "read: could not map logical address") {
-				dlog.Errorf(ctx, "dbg walk err: %v", err)
-			}
-		},
-	})
-
-	orphanedNodes := make(map[btrfsvol.LogicalAddr]int)
-	for _, devResults := range nodeScanResults {
-		for laddr := range devResults.FoundNodes {
-			if _, attached := visitedNodes[laddr]; !attached {
-				orphanedNodes[laddr] = 0
-			}
-		}
-	}
-	if len(visitedNodes)+len(orphanedNodes) != total {
-		panic("should not happen")
-	}
-	dlog.Infof(ctx,
-		"... (finished processing %v attached nodes, proceeding to process %v lost nodes, for a total of %v)",
-		len(visitedNodes), len(orphanedNodes), len(visitedNodes)+len(orphanedNodes))
-
-	// 'orphanedRoots' is a subset of 'orphanedNodes'; start with
-	// it as the complete orphanedNodes, and then remove entries.
-	orphanedRoots := make(map[btrfsvol.LogicalAddr]struct{}, len(orphanedNodes))
-	for node := range orphanedNodes {
-		orphanedRoots[node] = struct{}{}
-	}
-	for potentialRoot := range orphanedNodes {
-		if orphanedNodes[potentialRoot] > 1 {
-			continue
-		}
-		walkFromNode(ctx, fs, potentialRoot,
-			func(err *btrfstree.TreeError) {
-				// do nothing
-			},
-			btrfstree.TreeWalkHandler{
-				// Don't use `PreNode` because we don't want to run this on bad
-				// nodes (it'd screw up `len(visitedNodes)`).
-				Node: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
-					addr := path.Node(-1).ToNodeAddr
-					if addr != potentialRoot {
-						delete(orphanedRoots, addr)
-					}
-					if _, alreadyVisited := visitedNodes[addr]; alreadyVisited {
-						return iofs.SkipDir
-					}
-					visitedNodes[addr] = struct{}{}
-					progress(len(visitedNodes))
-					return nil
-				},
-			},
-		)
-	}
-
-	if len(visitedNodes) != total {
-		panic("should not happen")
-	}
-
-	dlog.Infof(ctx, "... identified %d lost+found nodes", len(orphanedRoots))
-	return orphanedRoots, nil
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reattach.go
index 32269fc..2b18aca 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reattach.go
@@ -18,7 +18,7 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
-func reAttachNodes(ctx context.Context, fs _FS, foundRoots map[btrfsvol.LogicalAddr]struct{}, rebuiltNodes map[btrfsvol.LogicalAddr]*RebuiltNode) error {
+func reAttachNodes(ctx context.Context, fs _FS, orphanedNodes map[btrfsvol.LogicalAddr]struct{}, rebuiltNodes map[btrfsvol.LogicalAddr]*RebuiltNode) error {
 	dlog.Info(ctx, "Attaching lost+found nodes to rebuilt nodes...")
 
 	sb, err := fs.Superblock()
@@ -48,19 +48,19 @@ func reAttachNodes(ctx context.Context, fs _FS, foundRoots map[btrfsvol.LogicalA
 	}
 	dlog.Info(ctx, "... done indexing")
 
-	// Attach foundRoots to the gaps.
+	// Attach orphanedNodes to the gaps.
 	dlog.Info(ctx, "... attaching nodes...")
 	lastPct := -1
 	progress := func(done int) {
-		pct := int(100 * float64(done) / float64(len(foundRoots)))
-		if pct != lastPct || done == len(foundRoots) {
+		pct := int(100 * float64(done) / float64(len(orphanedNodes)))
+		if pct != lastPct || done == len(orphanedNodes) {
 			dlog.Infof(ctx, "... %v%% (%v/%v)",
-				pct, done, len(foundRoots))
+				pct, done, len(orphanedNodes))
 			lastPct = pct
 		}
 	}
 	numAttached := 0
-	for i, foundLAddr := range maps.SortedKeys(foundRoots) {
+	for i, foundLAddr := range maps.SortedKeys(orphanedNodes) {
 		progress(i)
 		foundRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, foundLAddr, btrfstree.NodeExpectations{
 			LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: foundLAddr},
@@ -113,10 +113,10 @@ func reAttachNodes(ctx context.Context, fs _FS, foundRoots map[btrfsvol.LogicalA
 				foundRef.Addr)
 		}
 	}
-	progress(len(foundRoots))
+	progress(len(orphanedNodes))
 	dlog.Info(ctx, "... ... done attaching")
 
 	dlog.Infof(ctx, "... re-attached %d nodes (%v%% success rate)",
-		numAttached, int(100*float64(numAttached)/float64(len(foundRoots))))
+		numAttached, int(100*float64(numAttached)/float64(len(orphanedNodes))))
 	return nil
 }