1 files changed, 122 insertions, 0 deletions
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reattach.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reattach.go
new file mode 100644
index 0000000..2b18aca
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reattach.go
@@ -0,0 +1,122 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildnodes
+
+import (
+	"context"
+	"sort"
+
+	"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/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+)
+
+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()
+	if err != nil {
+		return err
+	}
+
+	// Index 'rebuiltNodes' for fast lookups.
+	dlog.Info(ctx, "... indexing rebuilt nodes...")
+	gaps := make(map[btrfsprim.ObjID]map[uint8][]*RebuiltNode)
+	maxLevel := make(map[btrfsprim.ObjID]uint8)
+	for _, node := range rebuiltNodes {
+		for treeID := range node.InTrees {
+			maxLevel[treeID] = slices.Max(maxLevel[treeID], node.Head.Level)
+			if gaps[treeID] == nil {
+				gaps[treeID] = make(map[uint8][]*RebuiltNode)
+			}
+			gaps[treeID][node.Head.Level] = append(gaps[treeID][node.Head.Level], node)
+		}
+	}
+	for _, byTreeID := range gaps {
+		for _, slice := range byTreeID {
+			sort.Slice(slice, func(i, j int) bool {
+				return slice[i].MinKey.Cmp(slice[j].MinKey) < 0
+			})
+		}
+	}
+	dlog.Info(ctx, "... done indexing")
+
+	// Attach orphanedNodes to the gaps.
+	dlog.Info(ctx, "... attaching nodes...")
+	lastPct := -1
+	progress := func(done int) {
+		pct := int(100 * float64(done) / float64(len(orphanedNodes)))
+		if pct != lastPct || done == len(orphanedNodes) {
+			dlog.Infof(ctx, "... %v%% (%v/%v)",
+				pct, done, len(orphanedNodes))
+			lastPct = pct
+		}
+	}
+	numAttached := 0
+	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},
+		})
+		if foundRef == nil {
+			return err
+		}
+		foundMinKey, ok := foundRef.Data.MinItem()
+		if !ok {
+			continue
+		}
+		foundMaxKey, ok := foundRef.Data.MaxItem()
+		if !ok {
+			continue
+		}
+		treeGaps := gaps[foundRef.Data.Head.Owner]
+		var attached bool
+		for level := foundRef.Data.Head.Level + 1; treeGaps != nil && level <= maxLevel[foundRef.Data.Head.Owner] && !attached; level++ {
+			parentGen, ok := treeGaps[level]
+			if !ok {
+				continue
+			}
+			parentIdx, ok := slices.Search(parentGen, func(parent *RebuiltNode) int {
+				switch {
+				case foundMinKey.Cmp(parent.MinKey) < 0:
+					// 'parent' is too far right
+					return -1
+				case foundMaxKey.Cmp(parent.MaxKey) > 0:
+					// 'parent' is too far left
+					return 1
+				default:
+					// just right
+					return 0
+				}
+			})
+			if !ok {
+				continue
+			}
+			parent := parentGen[parentIdx]
+			parent.BodyInternal = append(parent.BodyInternal, btrfstree.KeyPointer{
+				Key:        foundMinKey,
+				BlockPtr:   foundLAddr,
+				Generation: foundRef.Data.Head.Generation,
+			})
+			attached = true
+			numAttached++
+		}
+		if !attached {
+			dlog.Errorf(ctx, "could not find a broken node to attach node to reattach node@%v to",
+				foundRef.Addr)
+		}
+	}
+	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(orphanedNodes))))
+	return nil
+}