1 files changed, 166 insertions, 0 deletions
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
new file mode 100644
index 0000000..9fb5bf5
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
@@ -0,0 +1,166 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildnodes
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	iofs "io/fs"
+	"reflect"
+	"strings"
+
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"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/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+)
+
+// 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, nil, err
+	}
+	chunkTreeUUID, ok := getChunkTreeUUID(ctx, fs)
+	if !ok {
+		return nil, nil, fmt.Errorf("could not look up chunk tree UUID")
+	}
+
+	lastPct := -1
+	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)",
+				pct, done, total)
+			lastPct = pct
+		}
+	}
+
+	rebuiltNodes = make(map[btrfsvol.LogicalAddr]*RebuiltNode)
+	walkHandler := btrfstree.TreeWalkHandler{
+		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()
+			return nil
+		},
+		BadNode: func(path btrfstree.TreePath, _ *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
+			node := btrfstree.Node{
+				Size:         sb.NodeSize,
+				ChecksumType: sb.ChecksumType,
+				Head: btrfstree.NodeHeader{
+					MetadataUUID:  sb.EffectiveMetadataUUID(),
+					Addr:          path.Node(-1).ToNodeAddr,
+					ChunkTreeUUID: chunkTreeUUID,
+					//Owner:         TBD, // see RebuiltNode.InTrees
+					Generation: path.Node(-1).FromGeneration,
+					Level:      path.Node(-1).ToNodeLevel,
+				},
+			}
+			min, max := spanOfTreePath(fs, path)
+			if other, ok := rebuiltNodes[path.Node(-1).ToNodeAddr]; ok {
+				if !reflect.DeepEqual(other.Node, node) {
+					dlog.Errorf(ctx, "... mismatch: %v != %v", node, other.Node)
+					return err
+				}
+				if min.Cmp(other.MinKey) > 0 { // if min > other.MinKey {
+					other.MinKey = min // take the max of the two
+				}
+				if max.Cmp(other.MaxKey) < 0 { // if max < other.MaxKey {
+					other.MaxKey = max // take the min of the two
+				}
+				other.InTrees.Insert(path.Node(-1).FromTree)
+			} else {
+				rebuiltNodes[path.Node(-1).ToNodeAddr] = &RebuiltNode{
+					Err:     err.Error(),
+					MinKey:  min,
+					MaxKey:  max,
+					InTrees: containers.Set[btrfsprim.ObjID]{path.Node(-1).FromTree: struct{}{}},
+					Node:    node,
+				}
+			}
+			return err
+		},
+	}
+
+	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)
+			}
+		},
+	})
+
+	// 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,
+		)
+	}
+
+	if len(visitedNodes) != total {
+		panic("should not happen")
+	}
+
+	dlog.Infof(ctx, "... identified %d orphaned nodes and re-built %d nodes", len(orphanedNodes), len(rebuiltNodes))
+	return orphanedNodes, rebuiltNodes, nil
+}