btrfstree: Have LookupTreeRoot take a Context

5 files changed, 16 insertions, 13 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
index ada9f6f..3339270 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
@@ -59,7 +59,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, nodeList []btrfsvol.LogicalA
 	ret := ScanDevicesResult{
 		Superblock: *sb,
 
-		Graph: btrfsutil.NewGraph(*sb),
+		Graph: btrfsutil.NewGraph(ctx, *sb),
 
 		Flags: make(map[btrfsutil.ItemPtr]FlagsAndErr),
 		Names: make(map[btrfsutil.ItemPtr][]byte),
diff --git a/lib/btrfs/btrfstree/btree_forrest.go b/lib/btrfs/btrfstree/btree_forrest.go
index 0f46d42..d4317d6 100644
--- a/lib/btrfs/btrfstree/btree_forrest.go
+++ b/lib/btrfs/btrfstree/btree_forrest.go
@@ -5,6 +5,7 @@
 package btrfstree
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -24,7 +25,7 @@ type TreeRoot struct {
 
 // LookupTreeRoot is a utility function to help with implementing the
 // 'TreeOperator' interface.
-func LookupTreeRoot(fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
+func LookupTreeRoot(_ context.Context, fs TreeOperator, sb Superblock, treeID btrfsprim.ObjID) (*TreeRoot, error) {
 	switch treeID {
 	case btrfsprim.ROOT_TREE_OBJECTID:
 		return &TreeRoot{
diff --git a/lib/btrfs/btrfstree/btree_tree.go b/lib/btrfs/btrfstree/btree_tree.go
index 459f481..a34946e 100644
--- a/lib/btrfs/btrfstree/btree_tree.go
+++ b/lib/btrfs/btrfstree/btree_tree.go
@@ -28,7 +28,7 @@ func (fs TreeOperatorImpl) TreeWalk(ctx context.Context, treeID btrfsprim.ObjID,
 	if err != nil {
 		errHandle(&TreeError{Path: Path{{FromTree: treeID, ToMaxKey: btrfsprim.MaxKey}}, Err: err})
 	}
-	rootInfo, err := LookupTreeRoot(fs, *sb, treeID)
+	rootInfo, err := LookupTreeRoot(ctx, fs, *sb, treeID)
 	if err != nil {
 		errHandle(&TreeError{Path: Path{{FromTree: treeID, ToMaxKey: btrfsprim.MaxKey}}, Err: err})
 		return
@@ -405,11 +405,12 @@ func (fs TreeOperatorImpl) next(path Path, node *Node) (Path, *Node, error) {
 
 // TreeSearch implements the 'TreeOperator' interface.
 func (fs TreeOperatorImpl) TreeSearch(treeID btrfsprim.ObjID, searcher TreeSearcher) (Item, error) {
+	ctx := context.TODO()
 	sb, err := fs.Superblock()
 	if err != nil {
 		return Item{}, err
 	}
-	rootInfo, err := LookupTreeRoot(fs, *sb, treeID)
+	rootInfo, err := LookupTreeRoot(ctx, fs, *sb, treeID)
 	if err != nil {
 		return Item{}, err
 	}
@@ -430,11 +431,12 @@ func (fs TreeOperatorImpl) TreeLookup(treeID btrfsprim.ObjID, key btrfsprim.Key)
 
 // TreeSearchAll implements the 'TreeOperator' interface.
 func (fs TreeOperatorImpl) TreeSearchAll(treeID btrfsprim.ObjID, searcher TreeSearcher) ([]Item, error) {
+	ctx := context.TODO()
 	sb, err := fs.Superblock()
 	if err != nil {
 		return nil, err
 	}
-	rootInfo, err := LookupTreeRoot(fs, *sb, treeID)
+	rootInfo, err := LookupTreeRoot(ctx, fs, *sb, treeID)
 	if err != nil {
 		return nil, err
 	}
diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index 35848de..39e1cf2 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -111,8 +111,8 @@ func (g Graph) insertEdge(ptr *GraphEdge) {
 	g.EdgesTo[ptr.ToNode] = append(g.EdgesTo[ptr.ToNode], ptr)
 }
 
-func (g Graph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
-	treeInfo, err := btrfstree.LookupTreeRoot(nil, sb, treeID)
+func (g Graph) insertTreeRoot(ctx context.Context, sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
+	treeInfo, err := btrfstree.LookupTreeRoot(ctx, nil, sb, treeID)
 	if err != nil {
 		// This shouldn't ever happen for treeIDs that are
 		// mentioned directly in the superblock; which are the
@@ -131,7 +131,7 @@ func (g Graph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
 	})
 }
 
-func NewGraph(sb btrfstree.Superblock) Graph {
+func NewGraph(ctx context.Context, sb btrfstree.Superblock) Graph {
 	g := Graph{
 		Nodes:     make(map[btrfsvol.LogicalAddr]GraphNode),
 		BadNodes:  make(map[btrfsvol.LogicalAddr]error),
@@ -141,10 +141,10 @@ func NewGraph(sb btrfstree.Superblock) Graph {
 
 	// These 4 trees are mentioned directly in the superblock, so
 	// they are always seen.
-	g.insertTreeRoot(sb, btrfsprim.ROOT_TREE_OBJECTID)
-	g.insertTreeRoot(sb, btrfsprim.CHUNK_TREE_OBJECTID)
-	g.insertTreeRoot(sb, btrfsprim.TREE_LOG_OBJECTID)
-	g.insertTreeRoot(sb, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
+	g.insertTreeRoot(ctx, sb, btrfsprim.ROOT_TREE_OBJECTID)
+	g.insertTreeRoot(ctx, sb, btrfsprim.CHUNK_TREE_OBJECTID)
+	g.insertTreeRoot(ctx, sb, btrfsprim.TREE_LOG_OBJECTID)
+	g.insertTreeRoot(ctx, sb, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
 
 	return g
 }
diff --git a/lib/btrfsutil/old_rebuilt_forrest.go b/lib/btrfsutil/old_rebuilt_forrest.go
index abe3329..8ae7149 100644
--- a/lib/btrfsutil/old_rebuilt_forrest.go
+++ b/lib/btrfsutil/old_rebuilt_forrest.go
@@ -153,7 +153,7 @@ func (bt *OldRebuiltForrest) rawTreeWalk(treeID btrfsprim.ObjID, cacheEntry *old
 		cacheEntry.RootErr = err
 		return
 	}
-	root, err := btrfstree.LookupTreeRoot(bt, *sb, treeID)
+	root, err := btrfstree.LookupTreeRoot(bt.ctx, bt, *sb, treeID)
 	if err != nil {
 		cacheEntry.RootErr = err
 		return