Implement proper search for broken btrees

3 files changed, 269 insertions, 13 deletions

diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index 6665d09..e7c6cc2 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -19,6 +19,12 @@ import (
 )
 
 type Trees interface {
+	// TreeWalk walks a tree, triggering callbacks for every node,
+	// key-pointer, and item; as well as for any errors encountered.
+	//
+	// If the tree is valid, then everything is walked in key-order; but if
+	// the tree is broken, then ordering is not guaranteed.
+	//
 	// Canceling the Context causes TreeWalk to return early; no
 	// values from the Context are used.
 	//
@@ -178,43 +184,46 @@ func (e *TreeError) Error() string {
 	return fmt.Sprintf("%v: %v", e.Path, e.Err)
 }
 
-// A treeRoot is more-or-less a btrfsitem.Root, but simpler and generalized for
-type treeRoot struct {
+// A TreeRoot is more-or-less a btrfsitem.Root, but simpler; returned by
+// LookupTreeRoot.
+type TreeRoot struct {
 	TreeID     ObjID
 	RootNode   btrfsvol.LogicalAddr
 	Level      uint8
 	Generation Generation
 }
 
-func (fs *FS) lookupTree(treeID ObjID) (*treeRoot, error) {
+// LookupTreeRoot is a utility function to help with implementing the 'Trees'
+// interface.
+func LookupTreeRoot(fs Trees, treeID ObjID) (*TreeRoot, error) {
 	sb, err := fs.Superblock()
 	if err != nil {
 		return nil, err
 	}
 	switch treeID {
 	case ROOT_TREE_OBJECTID:
-		return &treeRoot{
+		return &TreeRoot{
 			TreeID:     treeID,
 			RootNode:   sb.RootTree,
 			Level:      sb.RootLevel,
 			Generation: sb.Generation, // XXX: same generation as LOG_TREE?
 		}, nil
 	case CHUNK_TREE_OBJECTID:
-		return &treeRoot{
+		return &TreeRoot{
 			TreeID:     treeID,
 			RootNode:   sb.ChunkTree,
 			Level:      sb.ChunkLevel,
 			Generation: sb.ChunkRootGeneration,
 		}, nil
 	case TREE_LOG_OBJECTID:
-		return &treeRoot{
+		return &TreeRoot{
 			TreeID:     treeID,
 			RootNode:   sb.LogTree,
 			Level:      sb.LogLevel,
 			Generation: sb.Generation, // XXX: same generation as ROOT_TREE?
 		}, nil
 	case BLOCK_GROUP_TREE_OBJECTID:
-		return &treeRoot{
+		return &TreeRoot{
 			TreeID:     treeID,
 			RootNode:   sb.BlockGroupRoot,
 			Level:      sb.BlockGroupRootLevel,
@@ -238,7 +247,7 @@ func (fs *FS) lookupTree(treeID ObjID) (*treeRoot, error) {
 		if !ok {
 			return nil, fmt.Errorf("malformed ROOT_ITEM for tree %v", treeID)
 		}
-		return &treeRoot{
+		return &TreeRoot{
 			TreeID:     treeID,
 			RootNode:   rootItemBody.ByteNr,
 			Level:      rootItemBody.Level,
@@ -265,7 +274,7 @@ func (fs *FS) TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeEr
 	path := TreePath{
 		TreeID: treeID,
 	}
-	rootInfo, err := fs.lookupTree(treeID)
+	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
 		errHandle(&TreeError{Path: path, Err: err})
 		return
@@ -278,6 +287,22 @@ func (fs *FS) TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeEr
 	fs.treeWalk(ctx, path, errHandle, cbs)
 }
 
+// TreeWalk is a utility function to help with implementing the 'Trees'
+// interface.
+func (fs *FS) RawTreeWalk(ctx context.Context, rootInfo TreeRoot, errHandle func(*TreeError), cbs TreeWalkHandler) {
+	path := TreePath{
+		TreeID: rootInfo.TreeID,
+		Nodes: []TreePathElem{
+			{
+				ItemIdx:   -1,
+				NodeAddr:  rootInfo.RootNode,
+				NodeLevel: rootInfo.Level,
+			},
+		},
+	}
+	fs.treeWalk(ctx, path, errHandle, cbs)
+}
+
 func (fs *FS) treeWalk(ctx context.Context, path TreePath, errHandle func(*TreeError), cbs TreeWalkHandler) {
 	if ctx.Err() != nil {
 		return
@@ -376,7 +401,7 @@ func (fs *FS) treeWalk(ctx context.Context, path TreePath, errHandle func(*TreeE
 	}
 }
 
-func (fs *FS) treeSearch(treeRoot treeRoot, fn func(Key) int) (TreePath, *util.Ref[btrfsvol.LogicalAddr, Node], error) {
+func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key) int) (TreePath, *util.Ref[btrfsvol.LogicalAddr, Node], error) {
 	path := TreePath{
 		TreeID: treeRoot.TreeID,
 		Nodes: []TreePathElem{
@@ -581,7 +606,7 @@ func (fs *FS) next(path TreePath, node *util.Ref[btrfsvol.LogicalAddr, Node]) (T
 }
 
 func (fs *FS) TreeSearch(treeID ObjID, fn func(Key) int) (Item, error) {
-	rootInfo, err := fs.lookupTree(treeID)
+	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
 		return Item{}, err
 	}
@@ -601,7 +626,7 @@ func (fs *FS) TreeLookup(treeID ObjID, key Key) (Item, error) {
 }
 
 func (fs *FS) TreeSearchAll(treeID ObjID, fn func(Key) int) ([]Item, error) {
-	rootInfo, err := fs.lookupTree(treeID)
+	rootInfo, err := LookupTreeRoot(fs, treeID)
 	if err != nil {
 		return nil, err
 	}
diff --git a/lib/btrfsprogs/btrfsinspect/mount.go b/lib/btrfsprogs/btrfsinspect/mount.go
index 9cb28dc..2df84f5 100644
--- a/lib/btrfsprogs/btrfsinspect/mount.go
+++ b/lib/btrfsprogs/btrfsinspect/mount.go
@@ -23,6 +23,7 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/linux"
 	"git.lukeshu.com/btrfs-progs-ng/lib/util"
 )
@@ -40,7 +41,7 @@ func MountRO(ctx context.Context, fs *btrfs.FS, mountpoint string) error {
 
 	rootSubvol := &subvolume{
 		Subvolume: btrfs.Subvolume{
-			FS:     fs,
+			FS:     btrfsutil.NewBrokenTrees(ctx, fs),
 			TreeID: btrfs.FS_TREE_OBJECTID,
 		},
 		DeviceName: deviceName,
diff --git a/lib/btrfsprogs/btrfsutil/broken_btree.go b/lib/btrfsprogs/btrfsutil/broken_btree.go
new file mode 100644
index 0000000..2835b35
--- /dev/null
+++ b/lib/btrfsprogs/btrfsutil/broken_btree.go
@@ -0,0 +1,230 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+import (
+	"context"
+	"fmt"
+	iofs "io/fs"
+	"sync"
+
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/rbtree"
+	"git.lukeshu.com/btrfs-progs-ng/lib/util"
+)
+
+type indexItem struct {
+	Key  btrfs.Key
+	Path btrfs.TreePath
+}
+
+type cachedIndex struct {
+	Index *rbtree.Tree[btrfs.Key, indexItem]
+	Err   error
+}
+
+type brokenTrees struct {
+	ctx   context.Context
+	inner *btrfs.FS
+
+	// btrfs.ROOT_TREE_OBJECTID
+	rootTreeMu    sync.Mutex
+	rootTreeIndex *cachedIndex
+	// for all other trees
+	treeMu      sync.Mutex
+	treeIndexes map[btrfs.ObjID]cachedIndex
+}
+
+var _ btrfs.Trees = (*brokenTrees)(nil)
+
+// NewBrokenTrees wraps a *btrfs.FS to support looking up information
+// from broken trees.
+//
+// Of the btrfs.FS.Tree{Verb}Trees methods:
+//
+//  - TreeWalk works on broken trees
+//  - TreeLookup relies on the tree being properly ordered (which a
+//    broken tree might not be).
+//  - TreeSearch relies on the tree being properly ordered (which a
+//    broken tree might not be).
+//  - TreeSearchAll relies on the tree being properly ordered (which a
+//    broken tree might not be), and a bad node may cause it to not
+//    return a truncated list of results.
+//
+// NewBrokenTrees attempts to remedy these deficiencies by using
+// .TreeWalk to build an out-of-FS index of all of the items in the
+// tree, and re-implements TreeLookup, TreeSearch, and TreeSearchAll
+// using that index.
+func NewBrokenTrees(ctx context.Context, inner *btrfs.FS) btrfs.Trees {
+	return &brokenTrees{
+		ctx:   ctx,
+		inner: inner,
+	}
+}
+
+func (bt *brokenTrees) treeIndex(treeID btrfs.ObjID) (*rbtree.Tree[btrfs.Key, indexItem], error) {
+	var treeRoot *btrfs.TreeRoot
+	var err error
+	if treeID == btrfs.ROOT_TREE_OBJECTID {
+		bt.rootTreeMu.Lock()
+		defer bt.rootTreeMu.Unlock()
+		if bt.rootTreeIndex != nil {
+			return bt.rootTreeIndex.Index, bt.rootTreeIndex.Err
+		}
+		treeRoot, err = btrfs.LookupTreeRoot(bt.inner, treeID)
+	} else {
+		bt.treeMu.Lock()
+		defer bt.treeMu.Unlock()
+		if bt.treeIndexes == nil {
+			bt.treeIndexes = make(map[btrfs.ObjID]cachedIndex)
+		}
+		if cacheEntry, exists := bt.treeIndexes[treeID]; exists {
+			return cacheEntry.Index, cacheEntry.Err
+		}
+		treeRoot, err = btrfs.LookupTreeRoot(bt, treeID)
+	}
+	var cacheEntry cachedIndex
+	if err != nil {
+		cacheEntry.Err = err
+	} else {
+		cacheEntry.Index = &rbtree.Tree[btrfs.Key, indexItem]{
+			KeyFn: func(item indexItem) btrfs.Key {
+				return item.Key
+			},
+		}
+		dlog.Infof(bt.ctx, "indexing tree %v...", treeID)
+		bt.inner.RawTreeWalk(
+			bt.ctx,
+			*treeRoot,
+			func(err *btrfs.TreeError) {
+				dlog.Error(bt.ctx, err)
+			},
+			btrfs.TreeWalkHandler{
+				Item: func(path btrfs.TreePath, item btrfs.Item) error {
+					if cacheEntry.Index.Lookup(item.Key) != nil {
+						panic(fmt.Errorf("dup key=%v in tree=%v", item.Key, treeID))
+					}
+					cacheEntry.Index.Insert(indexItem{
+						Path: path.DeepCopy(),
+						Key:  item.Key,
+					})
+					return nil
+				},
+			},
+		)
+		dlog.Infof(bt.ctx, "... done indexing tree %v", treeID)
+	}
+	if treeID == btrfs.ROOT_TREE_OBJECTID {
+		bt.rootTreeIndex = &cacheEntry
+	} else {
+		bt.treeIndexes[treeID] = cacheEntry
+	}
+	return cacheEntry.Index, cacheEntry.Err
+}
+
+func (bt *brokenTrees) TreeLookup(treeID btrfs.ObjID, key btrfs.Key) (btrfs.Item, error) {
+	item, err := bt.TreeSearch(treeID, key.Cmp)
+	if err != nil {
+		err = fmt.Errorf("item with key=%v: %w", key, err)
+	}
+	return item, err
+}
+
+func (bt *brokenTrees) TreeSearch(treeID btrfs.ObjID, fn func(btrfs.Key) int) (btrfs.Item, error) {
+	index, err := bt.treeIndex(treeID)
+	if err != nil {
+		return btrfs.Item{}, err
+	}
+	indexItem := index.Search(func(indexItem indexItem) int {
+		return fn(indexItem.Key)
+	})
+	if indexItem == nil {
+		return btrfs.Item{}, iofs.ErrNotExist
+	}
+
+	node, err := bt.inner.ReadNode(indexItem.Value.Path.Node(-2).NodeAddr)
+	if err != nil {
+		return btrfs.Item{}, err
+	}
+
+	item := node.Data.BodyLeaf[indexItem.Value.Path.Node(-1).ItemIdx]
+
+	return item, nil
+}
+
+func (bt *brokenTrees) TreeSearchAll(treeID btrfs.ObjID, fn func(btrfs.Key) int) ([]btrfs.Item, error) {
+	index, err := bt.treeIndex(treeID)
+	if err != nil {
+		return nil, err
+	}
+	indexItems := index.SearchRange(func(indexItem indexItem) int {
+		return fn(indexItem.Key)
+	})
+	if len(indexItems) == 0 {
+		return nil, iofs.ErrNotExist
+	}
+
+	ret := make([]btrfs.Item, len(indexItems))
+	var node *util.Ref[btrfsvol.LogicalAddr, btrfs.Node]
+	for i := range indexItems {
+		if node == nil || node.Addr != indexItems[i].Path.Node(-2).NodeAddr {
+			var err error
+			node, err = bt.inner.ReadNode(indexItems[i].Path.Node(-2).NodeAddr)
+			if err != nil {
+				return nil, err
+			}
+		}
+		ret[i] = node.Data.BodyLeaf[indexItems[i].Path.Node(-1).ItemIdx]
+	}
+	return ret, nil
+}
+
+func (bt *brokenTrees) TreeWalk(ctx context.Context, treeID btrfs.ObjID, errHandle func(*btrfs.TreeError), cbs btrfs.TreeWalkHandler) {
+	index, err := bt.treeIndex(treeID)
+	if err != nil {
+		errHandle(&btrfs.TreeError{
+			Path: btrfs.TreePath{
+				TreeID: treeID,
+			},
+			Err: err,
+		})
+		return
+	}
+	var node *util.Ref[btrfsvol.LogicalAddr, btrfs.Node]
+	_ = index.Walk(func(indexItem *rbtree.Node[indexItem]) error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		if bt.ctx.Err() != nil {
+			return bt.ctx.Err()
+		}
+		if cbs.Item != nil {
+			if node == nil || node.Addr != indexItem.Value.Path.Node(-2).NodeAddr {
+				var err error
+				node, err = bt.inner.ReadNode(indexItem.Value.Path.Node(-2).NodeAddr)
+				if err != nil {
+					errHandle(&btrfs.TreeError{Path: indexItem.Value.Path, Err: err})
+					return nil
+				}
+			}
+			item := node.Data.BodyLeaf[indexItem.Value.Path.Node(-1).ItemIdx]
+			if err := cbs.Item(indexItem.Value.Path, item); err != nil {
+				errHandle(&btrfs.TreeError{Path: indexItem.Value.Path, Err: err})
+			}
+		}
+		return nil
+	})
+}
+
+func (bt *brokenTrees) Superblock() (*btrfs.Superblock, error) {
+	return bt.inner.Superblock()
+}
+
+func (bt *brokenTrees) ReadAt(p []byte, off btrfsvol.LogicalAddr) (int, error) {
+	return bt.inner.ReadAt(p, off)
+}