wip rework rebuildnodes

2 files changed, 52 insertions, 33 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
index b29c201..38946b0 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
@@ -24,6 +24,22 @@ type uuidMap struct {
 	ObjID2UUID map[btrfsprim.ObjID]btrfsprim.UUID
 	UUID2ObjID map[btrfsprim.UUID]btrfsprim.ObjID
 	TreeParent map[btrfsprim.ObjID]btrfsprim.UUID
+
+	SeenTrees map[btrfsprim.ObjID]struct{}
+}
+
+func (m uuidMap) missingRootItems() map[btrfsprim.ObjID]struct{} {
+	missing := make(map[btrfsprim.ObjID]struct{})
+	for treeID := range m.SeenTrees {
+		if _, ok := m.ObjID2UUID[treeID]; !ok && treeID != btrfsprim.ROOT_TREE_OBJECTID {
+			missing[treeID] = struct{}{}
+			continue
+		}
+		if _, ok := m.TreeParent[treeID]; !ok && treeID >= btrfsprim.FIRST_FREE_OBJECTID && treeID <= btrfsprim.LAST_FREE_OBJECTID {
+			missing[treeID] = struct{}{}
+		}
+	}
+	return missing
 }
 
 func maybeSet[K, V comparable](name string, m map[K]V, k K, v V) error {
@@ -58,8 +74,9 @@ func buildUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sc
 		ObjID2UUID: make(map[btrfsprim.ObjID]btrfsprim.UUID),
 		UUID2ObjID: make(map[btrfsprim.UUID]btrfsprim.ObjID),
 		TreeParent: make(map[btrfsprim.ObjID]btrfsprim.UUID),
+
+		SeenTrees: make(map[btrfsprim.ObjID]struct{}),
 	}
-	seenTreeIDs := make(map[btrfsprim.ObjID]struct{})
 
 	progress()
 	for _, devResults := range scanResults {
@@ -94,7 +111,7 @@ func buildUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sc
 					}
 				}
 			}
-			seenTreeIDs[nodeRef.Data.Head.Owner] = struct{}{}
+			ret.SeenTrees[nodeRef.Data.Head.Owner] = struct{}{}
 			done++
 			progress()
 		}
@@ -104,16 +121,7 @@ func buildUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sc
 		panic("should not happen")
 	}
 
-	missing := make(map[btrfsprim.ObjID]struct{})
-	for treeID := range seenTreeIDs {
-		if _, ok := ret.ObjID2UUID[treeID]; !ok && treeID != btrfsprim.ROOT_TREE_OBJECTID {
-			missing[treeID] = struct{}{}
-			continue
-		}
-		if _, ok := ret.TreeParent[treeID]; !ok && treeID >= btrfsprim.FIRST_FREE_OBJECTID && treeID <= btrfsprim.LAST_FREE_OBJECTID {
-			missing[treeID] = struct{}{}
-		}
-	}
+	missing := ret.missingRootItems()
 	if len(missing) > 0 {
 		dlog.Errorf(ctx, "... could not find root items for %d trees: %v", len(missing), maps.SortedKeys(missing))
 	}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
index aa90582..e171e49 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
@@ -6,6 +6,7 @@ package rebuildnodes
 
 import (
 	"context"
+	"errors"
 	iofs "io/fs"
 
 	"github.com/datawire/dlib/dlog"
@@ -47,6 +48,32 @@ func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDev
 		}
 	}
 
+	var potentialRoot btrfsvol.LogicalAddr // zero for non-lost nodes, non-zero for lost nodes
+	nodeHandler := func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
+		if err != nil && (errors.Is(err, btrfstree.ErrNotANode) || errors.As(err, new(*btrfstree.IOError))) {
+			badNodes = append(badNodes, badNode{
+				Err:  err.Error(),
+				Path: path.DeepCopy(),
+			})
+			return err
+		}
+		addr := path.Node(-1).ToNodeAddr
+		visitedNodes[addr] = struct{}{}
+		if potentialRoot != 0 {
+			// lost node
+			if addr != potentialRoot {
+				delete(orphanedNodes, addr)
+			}
+			// TODO: Compare `nodeRef.Data.Head.Owner` and `path.Node(-1).FromTree` to
+			// maybe reconstruct a missing root item.  This is a sort of catch-22; we
+			// trust this data less because lost nodes may be discarded and not just
+			// lost, but non-lost nodes will never have a missing root item, so lost
+			// nodes are all we have to work with on this.
+		}
+		progress()
+		return nil
+	}
+
 	walkHandler := btrfstree.TreeWalkHandler{
 		PreNode: func(path btrfstree.TreePath) error {
 			addr := path.Node(-1).ToNodeAddr
@@ -57,18 +84,11 @@ func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDev
 			}
 			return nil
 		},
-		Node: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
-			addr := path.Node(-1).ToNodeAddr
-			visitedNodes[addr] = struct{}{}
-			progress()
-			return nil
+		Node: func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
+			return nodeHandler(path, nodeRef, nil)
 		},
-		BadNode: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
-			badNodes = append(badNodes, badNode{
-				Err:  err.Error(),
-				Path: path.DeepCopy(),
-			})
-			return err
+		BadNode: func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
+			return nodeHandler(path, nodeRef, err)
 		},
 	}
 
@@ -96,16 +116,7 @@ func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDev
 	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
-		}
+	for _, potentialRoot = range maps.SortedKeys(orphanedNodes) {
 		walkFromNode(ctx, fs, potentialRoot,
 			func(err *btrfstree.TreeError) {
 				// do nothing