wip

5 files changed, 225 insertions, 15 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
index cac9483..5afa783 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
@@ -27,22 +27,32 @@ import (
 )
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsvol.LogicalAddr]*RebuiltNode, error) {
+	uuidMap, err := buildTreeUUIDMap(ctx, fs, nodeScanResults)
+	if err != nil {
+		return nil, err
+	}
+
+	nfs := &RebuiltTrees{
+		inner:   fs,
+		uuidMap: uuidMap,
+	}
+
 	dlog.Info(ctx, "Identifying lost+found nodes...")
-	foundRoots, err := lostAndFoundNodes(ctx, fs, nodeScanResults)
+	foundRoots, err := lostAndFoundNodes(ctx, nfs, nodeScanResults)
 	if err != nil {
 		return nil, err
 	}
 	dlog.Infof(ctx, "... identified %d lost+found nodes", len(foundRoots))
 
 	dlog.Info(ctx, "Initializing nodes to re-build...")
-	rebuiltNodes, err := reInitBrokenNodes(ctx, fs, nodeScanResults, foundRoots)
+	rebuiltNodes, err := reInitBrokenNodes(ctx, nfs, nodeScanResults, foundRoots)
 	if err != nil {
 		return nil, err
 	}
 	dlog.Infof(ctx, "Initialized %d nodes", len(rebuiltNodes))
 
 	dlog.Info(ctx, "Attaching lost+found nodes to rebuilt nodes...")
-	if err := reAttachNodes(ctx, fs, foundRoots, rebuiltNodes); err != nil {
+	if err := reAttachNodes(ctx, nfs, foundRoots, rebuiltNodes); err != nil {
 		return nil, err
 	}
 	dlog.Info(ctx, "... done attaching")
@@ -68,7 +78,7 @@ func keyMm(key btrfsprim.Key) btrfsprim.Key {
 	return key
 }
 
-func spanOfTreePath(fs *btrfs.FS, path btrfstree.TreePath) (btrfsprim.Key, btrfsprim.Key) {
+func spanOfTreePath(fs _FS, path btrfstree.TreePath) (btrfsprim.Key, btrfsprim.Key) {
 	// tree root error
 	if len(path) == 0 {
 		return btrfsprim.Key{}, maxKey
@@ -100,7 +110,7 @@ func spanOfTreePath(fs *btrfs.FS, path btrfstree.TreePath) (btrfsprim.Key, btrfs
 	return low, high
 }
 
-func walkFromNode(ctx context.Context, fs *btrfs.FS, nodeAddr btrfsvol.LogicalAddr, errHandle func(*btrfstree.TreeError), cbs btrfstree.TreeWalkHandler) {
+func walkFromNode(ctx context.Context, fs _FS, nodeAddr btrfsvol.LogicalAddr, errHandle func(*btrfstree.TreeError), cbs btrfstree.TreeWalkHandler) {
 	sb, _ := fs.Superblock()
 	nodeRef, _ := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, nodeAddr, btrfstree.NodeExpectations{
 		LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: nodeAddr},
@@ -114,7 +124,7 @@ func walkFromNode(ctx context.Context, fs *btrfs.FS, nodeAddr btrfsvol.LogicalAd
 		Level:      nodeRef.Data.Head.Level,
 		Generation: nodeRef.Data.Head.Generation,
 	}
-	btrfstree.TreesImpl{NodeSource: fs}.RawTreeWalk(ctx, treeInfo, errHandle, cbs)
+	btrfstree.TreeOperatorImpl{NodeSource: fs}.RawTreeWalk(ctx, treeInfo, errHandle, cbs)
 }
 
 func countNodes(nodeScanResults btrfsinspect.ScanDevicesResult) int {
@@ -125,7 +135,7 @@ func countNodes(nodeScanResults btrfsinspect.ScanDevicesResult) int {
 	return cnt
 }
 
-func lostAndFoundNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsvol.LogicalAddr]struct{}, error) {
+func lostAndFoundNodes(ctx context.Context, fs _FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsvol.LogicalAddr]struct{}, error) {
 	lastPct := -1
 	total := countNodes(nodeScanResults)
 	progress := func(done int) {
@@ -210,7 +220,7 @@ func lostAndFoundNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsi
 	return orphanedRoots, nil
 }
 
-func getChunkTreeUUID(ctx context.Context, fs *btrfs.FS) (btrfsprim.UUID, bool) {
+func getChunkTreeUUID(ctx context.Context, fs _FS) (btrfsprim.UUID, bool) {
 	ctx, cancel := context.WithCancel(ctx)
 	var ret btrfsprim.UUID
 	var retOK bool
@@ -237,7 +247,7 @@ type RebuiltNode struct {
 	btrfstree.Node
 }
 
-func reInitBrokenNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult, foundRoots map[btrfsvol.LogicalAddr]struct{}) (map[btrfsvol.LogicalAddr]*RebuiltNode, error) {
+func reInitBrokenNodes(ctx context.Context, fs _FS, nodeScanResults btrfsinspect.ScanDevicesResult, foundRoots map[btrfsvol.LogicalAddr]struct{}) (map[btrfsvol.LogicalAddr]*RebuiltNode, error) {
 	sb, err := fs.Superblock()
 	if err != nil {
 		return nil, err
@@ -317,7 +327,7 @@ func reInitBrokenNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsi
 	return rebuiltNodes, nil
 }
 
-func reAttachNodes(ctx context.Context, fs *btrfs.FS, foundRoots map[btrfsvol.LogicalAddr]struct{}, rebuiltNodes map[btrfsvol.LogicalAddr]*RebuiltNode) error {
+func reAttachNodes(ctx context.Context, fs _FS, foundRoots map[btrfsvol.LogicalAddr]struct{}, rebuiltNodes map[btrfsvol.LogicalAddr]*RebuiltNode) error {
 	// Index 'rebuiltNodes' for fast lookups.
 	gaps := make(map[btrfsprim.ObjID]map[uint8][]*RebuiltNode)
 	maxLevel := make(map[btrfsprim.ObjID]uint8)
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go
new file mode 100644
index 0000000..17116c6
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go
@@ -0,0 +1,94 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildnodes
+
+import (
+	"context"
+	"fmt"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"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/diskio"
+)
+
+type RebuiltTrees struct {
+	inner   *btrfs.FS
+	uuidMap treeUUIDMap
+	nodes   map[btrfsvol.LogicalAddr]*RebuiltNode
+}
+
+type _FS interface {
+	diskio.File[btrfsvol.LogicalAddr]
+	ParentTree(btrfsprim.ObjID) (btrfsprim.ObjID, bool)
+	btrfstree.NodeSource
+	btrfstree.TreeOperator
+}
+
+// diskio.File
+
+func (fs *RebuiltTrees) Name() string               { return fs.inner.Name() }
+func (fs *RebuiltTrees) Size() btrfsvol.LogicalAddr { return fs.inner.Size() }
+func (fs *RebuiltTrees) Close() error               { return fs.inner.Close() }
+func (fs *RebuiltTrees) ReadAt(p []byte, off btrfsvol.LogicalAddr) (int, error) {
+	sb, err := fs.Superblock()
+	if err != nil {
+		return 0, err
+	}
+	if rebuilt, ok := fs.nodes[off]; ok && len(p) == int(sb.NodeSize) {
+		bs, err := rebuilt.Node.MarshalBinary()
+		if err != nil {
+			panic(fmt.Errorf("should not happen: %w", err))
+		}
+		if len(p) != len(bs) {
+			panic(fmt.Errorf("should not happen: sb.NodeSize=%v but node marshaled to %v", sb.NodeSize, len(bs)))
+		}
+		return copy(p, bs), nil
+	}
+	return fs.inner.ReadAt(p, off)
+}
+func (fs *RebuiltTrees) WriteAt(p []byte, off btrfsvol.LogicalAddr) (int, error) {
+	return fs.inner.WriteAt(p, off)
+}
+
+// btrfstree.NodeSource backend
+
+func (fs *RebuiltTrees) ParentTree(tree btrfsprim.ObjID) (btrfsprim.ObjID, bool) {
+	if tree < btrfsprim.FIRST_FREE_OBJECTID {
+		return 0, false
+	}
+	parentUUID := fs.uuidMap.TreeParent[tree]
+	if parentUUID == (btrfsprim.UUID{}) {
+		return 0, false
+	}
+	parentObjID, ok := fs.uuidMap.UUID2ObjID[parentUUID]
+	if !ok {
+		panic(fmt.Errorf("should not happen: could not resolve parentUUID=%v to an ObjID", parentUUID))
+	}
+	return parentObjID, true
+}
+
+// btrfstree.NodeSource
+
+func (fs *RebuiltTrees) Superblock() (*btrfstree.Superblock, error) { return fs.inner.Superblock() }
+func (fs *RebuiltTrees) ReadNode(path btrfstree.TreePath) (*diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], error) {
+	return btrfstree.FSReadNode(fs, path)
+}
+
+// btrfstree.TreeOperator
+
+func (fs *RebuiltTrees) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID, errHandle func(*btrfstree.TreeError), cbs btrfstree.TreeWalkHandler) {
+	btrfstree.TreeOperatorImpl{NodeSource: fs}.TreeWalk(ctx, treeID, errHandle, cbs)
+}
+func (fs *RebuiltTrees) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key) (btrfstree.Item, error) {
+	return btrfstree.TreeOperatorImpl{NodeSource: fs}.TreeLookup(treeID, key)
+}
+func (fs *RebuiltTrees) TreeSearch(treeID btrfsprim.ObjID, fn func(key btrfsprim.Key, size uint32) int) (btrfstree.Item, error) {
+	return btrfstree.TreeOperatorImpl{NodeSource: fs}.TreeSearch(treeID, fn)
+}
+func (fs *RebuiltTrees) TreeSearchAll(treeID btrfsprim.ObjID, fn func(key btrfsprim.Key, size uint32) int) ([]btrfstree.Item, error) {
+	return btrfstree.TreeOperatorImpl{NodeSource: fs}.TreeSearchAll(treeID, fn)
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap.go
new file mode 100644
index 0000000..0179a6e
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap.go
@@ -0,0 +1,107 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildnodes
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
+	"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/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+)
+
+type treeUUIDMap struct {
+	ObjID2UUID map[btrfsprim.ObjID]btrfsprim.UUID
+	UUID2ObjID map[btrfsprim.UUID]btrfsprim.ObjID
+	TreeParent map[btrfsprim.ObjID]btrfsprim.UUID
+}
+
+func maybeSet[K, V comparable](m map[K]V, k K, v V) error {
+	if other, conflict := m[k]; conflict && other != v {
+		return fmt.Errorf("conflict: %v can't have both %v and %v", k, other, v)
+	}
+	m[k] = v
+	return nil
+}
+
+func buildTreeUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (treeUUIDMap, error) {
+	dlog.Infof(ctx, "Building table of tree ObjID←→UUID...")
+
+	sb, err := fs.Superblock()
+	if err != nil {
+		return treeUUIDMap{}, nil
+	}
+
+	lastPct := -1
+	total := countNodes(scanResults)
+	done := 0
+	progress := func() {
+		pct := int(100 * float64(done) / float64(total))
+		if pct != lastPct || done == total {
+			dlog.Infof(ctx, "... %v%% (%v/%v)",
+				pct, done, total)
+			lastPct = pct
+		}
+	}
+
+	ret := treeUUIDMap{
+		ObjID2UUID: make(map[btrfsprim.ObjID]btrfsprim.UUID),
+		UUID2ObjID: make(map[btrfsprim.UUID]btrfsprim.ObjID),
+		TreeParent: make(map[btrfsprim.ObjID]btrfsprim.UUID),
+	}
+	treeIDs := make(map[btrfsprim.ObjID]struct{})
+
+	progress()
+	for _, devResults := range scanResults {
+		for laddr := range devResults.FoundNodes {
+			nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{
+				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
+			})
+			if err != nil {
+				return treeUUIDMap{}, nil
+			}
+			for _, item := range nodeRef.Data.BodyLeaf {
+				itemBody, ok := item.Body.(btrfsitem.Root)
+				if !ok {
+					continue
+				}
+				if err := maybeSet(ret.ObjID2UUID, item.Key.ObjectID, itemBody.UUID); err != nil {
+					return treeUUIDMap{}, err
+				}
+				if err := maybeSet(ret.TreeParent, item.Key.ObjectID, itemBody.ParentUUID); err != nil {
+					return treeUUIDMap{}, err
+				}
+				if err := maybeSet(ret.UUID2ObjID, itemBody.UUID, item.Key.ObjectID); err != nil {
+					return treeUUIDMap{}, err
+				}
+			}
+			treeIDs[nodeRef.Data.Head.Owner] = struct{}{}
+			done++
+			progress()
+		}
+	}
+	progress()
+
+	missing := make(map[btrfsprim.ObjID]struct{})
+	for treeID := range treeIDs {
+		if _, ok := ret.ObjID2UUID[treeID]; !ok {
+			missing[treeID] = struct{}{}
+		}
+	}
+	if len(missing) > 0 {
+		return treeUUIDMap{}, fmt.Errorf("could not find root items for trees %v", maps.SortedKeys(missing))
+	}
+
+	dlog.Info(ctx, ".. done building table")
+	return ret, nil
+}
diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go
index 6afaceb..7341d94 100644
--- a/lib/btrfsprogs/btrfsutil/broken_btree.go
+++ b/lib/btrfsprogs/btrfsutil/broken_btree.go
@@ -102,7 +102,7 @@ type brokenTrees struct {
 	treeIndexes map[btrfsprim.ObjID]cachedIndex
 }
 
-var _ btrfstree.Trees = (*brokenTrees)(nil)
+var _ btrfstree.TreeOperator = (*brokenTrees)(nil)
 
 // NewBrokenTrees wraps a *btrfs.FS to support looking up information
 // from broken trees.
@@ -123,7 +123,7 @@ var _ btrfstree.Trees = (*brokenTrees)(nil)
 // tree, and re-implements TreeLookup, TreeSearch, and TreeSearchAll
 // using that index.
 func NewBrokenTrees(ctx context.Context, inner *btrfs.FS) interface {
-	btrfstree.Trees
+	btrfstree.TreeOperator
 	Superblock() (*btrfstree.Superblock, error)
 	ReadAt(p []byte, off btrfsvol.LogicalAddr) (int, error)
 } {
@@ -171,7 +171,7 @@ func (bt *brokenTrees) treeIndex(treeID btrfsprim.ObjID) cachedIndex {
 			KeyFn: func(iv indexValue) btrfsprim.Key { return iv.Key },
 		}
 		dlog.Infof(bt.ctx, "indexing tree %v...", treeID)
-		btrfstree.TreesImpl{NodeSource: bt.inner}.RawTreeWalk(
+		btrfstree.TreeOperatorImpl{NodeSource: bt.inner}.RawTreeWalk(
 			bt.ctx,
 			*treeRoot,
 			func(err *btrfstree.TreeError) {
diff --git a/lib/btrfsprogs/btrfsutil/walk.go b/lib/btrfsprogs/btrfsutil/walk.go
index 2809737..355976a 100644
--- a/lib/btrfsprogs/btrfsutil/walk.go
+++ b/lib/btrfsprogs/btrfsutil/walk.go
@@ -8,7 +8,6 @@ import (
 	"context"
 	"fmt"
 
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
@@ -37,7 +36,7 @@ type WalkAllTreesHandler struct {
 // WalkAllTrees walks all trees in a *btrfs.FS.  Rather than returning
 // an error, it calls errCb each time an error is encountered.  The
 // error will always be of type WalkError.
-func WalkAllTrees(ctx context.Context, fs *btrfs.FS, cbs WalkAllTreesHandler) {
+func WalkAllTrees(ctx context.Context, fs btrfstree.TreeOperator, cbs WalkAllTreesHandler) {
 	var treeName string
 
 	trees := []struct {