Have TreeWalk be cancelable

4 files changed, 10 insertions, 9 deletions

diff --git a/cmd/btrfs-fsck/main.go b/cmd/btrfs-fsck/main.go
index fd9eb1b..e4e2b25 100644
--- a/cmd/btrfs-fsck/main.go
+++ b/cmd/btrfs-fsck/main.go
@@ -37,7 +37,7 @@ func Main(ctx context.Context, imgfilenames ...string) (err error) {
 		return err
 	}
 
-	pass2(fs, foundNodes)
+	pass2(ctx, fs, foundNodes)
 
 	return nil
 }
diff --git a/cmd/btrfs-fsck/pass1.go b/cmd/btrfs-fsck/pass1.go
index e41694a..a625193 100644
--- a/cmd/btrfs-fsck/pass1.go
+++ b/cmd/btrfs-fsck/pass1.go
@@ -22,7 +22,7 @@ func pass1(ctx context.Context, fs *btrfs.FS, superblock *btrfs.Superblock) (map
 
 	fmt.Printf("Pass 1: ... walking fs\n")
 	visitedNodes := make(map[btrfsvol.LogicalAddr]struct{})
-	btrfsutil.WalkAllTrees(fs, btrfsutil.WalkAllTreesHandler{
+	btrfsutil.WalkAllTrees(ctx, fs, btrfsutil.WalkAllTreesHandler{
 		TreeWalkHandler: btrfs.TreeWalkHandler{
 			Node: func(path btrfs.TreePath, node *util.Ref[btrfsvol.LogicalAddr, btrfs.Node]) error {
 				visitedNodes[node.Addr] = struct{}{}
@@ -57,12 +57,12 @@ func pass1(ctx context.Context, fs *btrfs.FS, superblock *btrfs.Superblock) (map
 	btrfsinspect.PrintPhysicalSpace(os.Stdout, fs)
 
 	fmt.Printf("Pass 1: ... writing re-constructed chunks\n")
-	pass1WriteReconstructedChunks(fs)
+	pass1WriteReconstructedChunks(ctx, fs)
 
 	return fsFoundNodes, nil
 }
 
-func pass1WriteReconstructedChunks(fs *btrfs.FS) {
+func pass1WriteReconstructedChunks(ctx context.Context, fs *btrfs.FS) {
 	superblock, _ := fs.Superblock()
 
 	// FIXME(lukeshu): OK, so this just assumes that all the
@@ -145,7 +145,7 @@ func pass1WriteReconstructedChunks(fs *btrfs.FS) {
 		fmt.Printf("Pass 1: ... write new node: error: %v\n", err)
 	}
 
-	if err := fs.ReInit(); err != nil {
+	if err := fs.ReInit(ctx); err != nil {
 		fmt.Printf("Pass 1: ... re-init mappings: %v\n", err)
 	}
 
diff --git a/cmd/btrfs-fsck/pass2.go b/cmd/btrfs-fsck/pass2.go
index 78bd2e5..ef55d93 100644
--- a/cmd/btrfs-fsck/pass2.go
+++ b/cmd/btrfs-fsck/pass2.go
@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
@@ -13,11 +14,11 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/util"
 )
 
-func pass2(fs *btrfs.FS, foundNodes map[btrfsvol.LogicalAddr]struct{}) {
+func pass2(ctx context.Context, fs *btrfs.FS, foundNodes map[btrfsvol.LogicalAddr]struct{}) {
 	fmt.Printf("\nPass 2: orphaned nodes\n")
 
 	visitedNodes := make(map[btrfsvol.LogicalAddr]struct{})
-	btrfsutil.WalkAllTrees(fs, btrfsutil.WalkAllTreesHandler{
+	btrfsutil.WalkAllTrees(ctx, fs, btrfsutil.WalkAllTreesHandler{
 		TreeWalkHandler: btrfs.TreeWalkHandler{
 			Node: func(path btrfs.TreePath, node *util.Ref[btrfsvol.LogicalAddr, btrfs.Node]) error {
 				visitedNodes[node.Addr] = struct{}{}
diff --git a/cmd/btrfs-rec/inspect_lstrees.go b/cmd/btrfs-rec/inspect_lstrees.go
index 0a98369..590ec23 100644
--- a/cmd/btrfs-rec/inspect_lstrees.go
+++ b/cmd/btrfs-rec/inspect_lstrees.go
@@ -26,10 +26,10 @@ func init() {
 			Short: "A brief view what types of items are in each tree",
 			Args:  cliutil.WrapPositionalArgs(cobra.NoArgs),
 		},
-		RunE: func(fs *btrfs.FS, _ *cobra.Command, _ []string) error {
+		RunE: func(fs *btrfs.FS, cmd *cobra.Command, _ []string) error {
 			var treeErrCnt int
 			var treeItemCnt map[btrfsitem.Type]int
-			btrfsutil.WalkAllTrees(fs, btrfsutil.WalkAllTreesHandler{
+			btrfsutil.WalkAllTrees(cmd.Context(), fs, btrfsutil.WalkAllTreesHandler{
 				PreTree: func(name string, treeID btrfs.ObjID) {
 					treeErrCnt = 0
 					treeItemCnt = make(map[btrfsitem.Type]int)