Have TreeWalk be cancelable

5 files changed, 47 insertions, 12 deletions

diff --git a/lib/btrfs/io2_lv.go b/lib/btrfs/io2_lv.go
index eee58d5..b3f9276 100644
--- a/lib/btrfs/io2_lv.go
+++ b/lib/btrfs/io2_lv.go
@@ -38,7 +38,7 @@ func (fs *FS) AddDevice(ctx context.Context, dev *Device) error {
 	}
 	fs.cacheSuperblocks = nil
 	fs.cacheSuperblock = nil
-	if err := fs.initDev(*sb); err != nil {
+	if err := fs.initDev(ctx, *sb); err != nil {
 		dlog.Errorf(ctx, "error: AddDevice: %q: %v", dev.Name(), err)
 	}
 	return nil
@@ -133,21 +133,21 @@ func (fs *FS) Superblock() (*Superblock, error) {
 	return &sbs[0].Data, nil
 }
 
-func (fs *FS) ReInit() error {
+func (fs *FS) ReInit(ctx context.Context) error {
 	fs.LV.ClearMappings()
 	for _, dev := range fs.LV.PhysicalVolumes() {
 		sb, err := dev.Superblock()
 		if err != nil {
 			return fmt.Errorf("file %q: %w", dev.Name(), err)
 		}
-		if err := fs.initDev(*sb); err != nil {
+		if err := fs.initDev(ctx, *sb); err != nil {
 			return fmt.Errorf("file %q: %w", dev.Name(), err)
 		}
 	}
 	return nil
 }
 
-func (fs *FS) initDev(sb Superblock) error {
+func (fs *FS) initDev(ctx context.Context, sb Superblock) error {
 	syschunks, err := sb.ParseSysChunkArray()
 	if err != nil {
 		return err
@@ -160,7 +160,7 @@ func (fs *FS) initDev(sb Superblock) error {
 		}
 	}
 	var errs derror.MultiError
-	fs.TreeWalk(CHUNK_TREE_OBJECTID,
+	fs.TreeWalk(ctx, CHUNK_TREE_OBJECTID,
 		func(err *TreeError) {
 			errs = append(errs, err)
 		},
diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index 1cb1d74..5bbb8d1 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -5,6 +5,7 @@
 package btrfs
 
 import (
+	"context"
 	"fmt"
 	"io"
 	iofs "io/fs"
@@ -18,6 +19,9 @@ import (
 )
 
 type Trees interface {
+	// Canceling the Context causes TreeWalk to return early; no
+	// values from the Context are used.
+	//
 	// The lifecycle of callbacks is:
 	//
 	//     001 .PreNode()
@@ -31,7 +35,7 @@ type Trees interface {
 	//           else:
 	//     004     .Item() (or .BadItem())
 	//     007 .PostNode()
-	TreeWalk(treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHandler)
+	TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHandler)
 
 	TreeLookup(treeID ObjID, key Key) (Item, error)
 	TreeSearch(treeID ObjID, fn func(Key) int) (Item, error)
@@ -247,7 +251,7 @@ type TreeWalkHandler struct {
 	BadItem func(TreePath, Item) error
 }
 
-func (fs *FS) TreeWalk(treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
+func (fs *FS) TreeWalk(ctx context.Context, treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHandler) {
 	path := TreePath{
 		TreeID: treeID,
 	}
@@ -261,10 +265,13 @@ func (fs *FS) TreeWalk(treeID ObjID, errHandle func(*TreeError), cbs TreeWalkHan
 		NodeAddr:  rootInfo.RootNode,
 		NodeLevel: rootInfo.Level,
 	})
-	fs.treeWalk(path, errHandle, cbs)
+	fs.treeWalk(ctx, path, errHandle, cbs)
 }
 
-func (fs *FS) treeWalk(path TreePath, errHandle func(*TreeError), cbs TreeWalkHandler) {
+func (fs *FS) treeWalk(ctx context.Context, path TreePath, errHandle func(*TreeError), cbs TreeWalkHandler) {
+	if ctx.Err() != nil {
+		return
+	}
 	if path.Nodes[len(path.Nodes)-1].NodeAddr == 0 {
 		return
 	}
@@ -273,8 +280,14 @@ func (fs *FS) treeWalk(path TreePath, errHandle func(*TreeError), cbs TreeWalkHa
 		if err := cbs.PreNode(path); err != nil {
 			errHandle(&TreeError{Path: path, Err: err})
 		}
+		if ctx.Err() != nil {
+			return
+		}
 	}
 	node, err := fs.readNodeAtLevel(path.Nodes[len(path.Nodes)-1].NodeAddr, path.Nodes[len(path.Nodes)-1].NodeLevel)
+	if ctx.Err() != nil {
+		return
+	}
 	if err != nil && node != nil && cbs.BadNode != nil {
 		// opportunity to fix the node
 		err = cbs.BadNode(path, node, err)
@@ -288,6 +301,9 @@ func (fs *FS) treeWalk(path TreePath, errHandle func(*TreeError), cbs TreeWalkHa
 			}
 		}
 	}
+	if ctx.Err() != nil {
+		return
+	}
 	if node != nil {
 		for i, item := range node.Data.BodyInternal {
 			itemPath := path.Append(TreePathElem{
@@ -299,12 +315,18 @@ func (fs *FS) treeWalk(path TreePath, errHandle func(*TreeError), cbs TreeWalkHa
 				if err := cbs.PreKeyPointer(itemPath, item); err != nil {
 					errHandle(&TreeError{Path: itemPath, Err: err})
 				}
+				if ctx.Err() != nil {
+					return
+				}
 			}
-			fs.treeWalk(itemPath, errHandle, cbs)
+			fs.treeWalk(ctx, itemPath, errHandle, cbs)
 			if cbs.PostKeyPointer != nil {
 				if err := cbs.PostKeyPointer(itemPath, item); err != nil {
 					errHandle(&TreeError{Path: itemPath, Err: err})
 				}
+				if ctx.Err() != nil {
+					return
+				}
 			}
 		}
 		for i, item := range node.Data.BodyLeaf {
@@ -318,12 +340,18 @@ func (fs *FS) treeWalk(path TreePath, errHandle func(*TreeError), cbs TreeWalkHa
 					if err := cbs.BadItem(itemPath, item); err != nil {
 						errHandle(&TreeError{Path: itemPath, Err: err})
 					}
+					if ctx.Err() != nil {
+						return
+					}
 				}
 			} else {
 				if cbs.Item != nil {
 					if err := cbs.Item(itemPath, item); err != nil {
 						errHandle(&TreeError{Path: itemPath, Err: err})
 					}
+					if ctx.Err() != nil {
+						return
+					}
 				}
 			}
 		}
@@ -332,6 +360,9 @@ func (fs *FS) treeWalk(path TreePath, errHandle func(*TreeError), cbs TreeWalkHa
 		if err := cbs.PostNode(path, node); err != nil {
 			errHandle(&TreeError{Path: path, Err: err})
 		}
+		if ctx.Err() != nil {
+			return
+		}
 	}
 }
 
diff --git a/lib/btrfsprogs/btrfsinspect/print_tree.go b/lib/btrfsprogs/btrfsinspect/print_tree.go
index 6c4f550..142f5e5 100644
--- a/lib/btrfsprogs/btrfsinspect/print_tree.go
+++ b/lib/btrfsprogs/btrfsinspect/print_tree.go
@@ -44,6 +44,7 @@ func DumpTrees(ctx context.Context, out io.Writer, fs *btrfs.FS) {
 		printTree(ctx, out, fs, btrfs.BLOCK_GROUP_TREE_OBJECTID)
 	}
 	fs.TreeWalk(
+		ctx,
 		btrfs.ROOT_TREE_OBJECTID,
 		func(err *btrfs.TreeError) {
 			dlog.Error(ctx, err)
@@ -343,6 +344,7 @@ func printTree(ctx context.Context, out io.Writer, fs *btrfs.FS, treeID btrfs.Ob
 	}
 	handlers.BadItem = handlers.Item
 	fs.TreeWalk(
+		ctx,
 		treeID,
 		func(err *btrfs.TreeError) {
 			dlog.Error(ctx, err)
diff --git a/lib/btrfsprogs/btrfsrepair/clearnodes.go b/lib/btrfsprogs/btrfsrepair/clearnodes.go
index dbe1c3a..b38a333 100644
--- a/lib/btrfsprogs/btrfsrepair/clearnodes.go
+++ b/lib/btrfsprogs/btrfsrepair/clearnodes.go
@@ -21,7 +21,7 @@ func ClearBadNodes(ctx context.Context, fs *btrfs.FS) error {
 	var uuidsInited bool
 	var metadataUUID, chunkTreeUUID btrfs.UUID
 
-	btrfsutil.WalkAllTrees(fs, btrfsutil.WalkAllTreesHandler{
+	btrfsutil.WalkAllTrees(ctx, fs, btrfsutil.WalkAllTreesHandler{
 		Err: func(err *btrfsutil.WalkError) {
 			dlog.Error(ctx, err)
 		},
diff --git a/lib/btrfsprogs/btrfsutil/walk.go b/lib/btrfsprogs/btrfsutil/walk.go
index d2b6367..c597e5e 100644
--- a/lib/btrfsprogs/btrfsutil/walk.go
+++ b/lib/btrfsprogs/btrfsutil/walk.go
@@ -5,6 +5,7 @@
 package btrfsutil
 
 import (
+	"context"
 	"fmt"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
@@ -34,7 +35,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(fs *btrfs.FS, cbs WalkAllTreesHandler) {
+func WalkAllTrees(ctx context.Context, fs *btrfs.FS, cbs WalkAllTreesHandler) {
 	var treeName string
 
 	trees := []struct {
@@ -83,6 +84,7 @@ func WalkAllTrees(fs *btrfs.FS, cbs WalkAllTreesHandler) {
 			cbs.PreTree(treeName, tree.ID)
 		}
 		fs.TreeWalk(
+			ctx,
 			tree.ID,
 			func(err *btrfs.TreeError) { cbs.Err(&WalkError{TreeName: treeName, Err: err}) },
 			cbs.TreeWalkHandler,