Factor out the device-scan loop from rebuildmappings in to btrfsutil

2 files changed, 263 insertions, 179 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildmappings/scan.go b/cmd/btrfs-rec/inspect/rebuildmappings/scan.go
index 2128a48..f9c3bf3 100644
--- a/cmd/btrfs-rec/inspect/rebuildmappings/scan.go
+++ b/cmd/btrfs-rec/inspect/rebuildmappings/scan.go
@@ -6,55 +6,26 @@ package rebuildmappings
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"strings"
-	"sync"
-	"time"
 
-	"github.com/datawire/dlib/dgroup"
 	"github.com/datawire/dlib/dlog"
 
-	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct"
 	"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/btrfssum"
 	"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/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
-type ScanDevicesResult map[btrfsvol.DeviceID]ScanOneDeviceResult
+// Result types ////////////////////////////////////////////////////////////////
 
-func ScanDevices(ctx context.Context, fs *btrfs.FS) (ScanDevicesResult, error) {
-	grp := dgroup.NewGroup(ctx, dgroup.GroupConfig{})
-	var mu sync.Mutex
-	result := make(map[btrfsvol.DeviceID]ScanOneDeviceResult)
-	for id, dev := range fs.LV.PhysicalVolumes() {
-		id := id
-		dev := dev
-		grp.Go(fmt.Sprintf("dev-%d", id), func(ctx context.Context) error {
-			sb, err := dev.Superblock()
-			if err != nil {
-				return err
-			}
-			devResult, err := ScanOneDevice(ctx, dev, *sb)
-			if err != nil {
-				return err
-			}
-			mu.Lock()
-			result[id] = devResult
-			mu.Unlock()
-			return nil
-		})
-	}
-	if err := grp.Wait(); err != nil {
-		return nil, err
-	}
-	return result, nil
-}
+type ScanDevicesResult = map[btrfsvol.DeviceID]ScanOneDeviceResult
 
 type ScanOneDeviceResult struct {
 	Checksums        btrfssum.SumRun[btrfsvol.PhysicalAddr]
@@ -85,9 +56,27 @@ func (a SysExtentCSum) Compare(b SysExtentCSum) int {
 	return containers.NativeCompare(a.Sums.Addr, b.Sums.Addr)
 }
 
-type scanStats struct {
-	textui.Portion[btrfsvol.PhysicalAddr]
+// Convenience functions for those types ///////////////////////////////////////
+
+func ScanDevices(ctx context.Context, fs *btrfs.FS) (ScanDevicesResult, error) {
+	return btrfsutil.ScanDevices[scanStats, ScanOneDeviceResult](ctx, fs, newDeviceScanner)
+}
 
+// ScanOneDevice mostly mimics btrfs-progs
+// cmds/rescue-chunk-recover.c:scan_one_device().
+func ScanOneDevice(ctx context.Context, dev *btrfs.Device) (ScanOneDeviceResult, error) {
+	return btrfsutil.ScanOneDevice[scanStats, ScanOneDeviceResult](ctx, dev, newDeviceScanner)
+}
+
+// scanner implementation //////////////////////////////////////////////////////
+
+type deviceScanner struct {
+	alg    btrfssum.CSumType
+	sums   strings.Builder
+	result ScanOneDeviceResult
+}
+
+type scanStats struct {
 	NumFoundNodes       int
 	NumFoundChunks      int
 	NumFoundBlockGroups int
@@ -96,8 +85,7 @@ type scanStats struct {
 }
 
 func (s scanStats) String() string {
-	return textui.Sprintf("scanned %v (found: %v nodes, %v chunks, %v block groups, %v dev extents, %v sum items)",
-		s.Portion,
+	return textui.Sprintf("found: %v nodes, %v chunks, %v block groups, %v dev extents, %v sum items",
 		s.NumFoundNodes,
 		s.NumFoundChunks,
 		s.NumFoundBlockGroups,
@@ -105,156 +93,104 @@ func (s scanStats) String() string {
 		s.NumFoundExtentCSums)
 }
 
-var sbSize = btrfsvol.PhysicalAddr(binstruct.StaticSize(btrfstree.Superblock{}))
-
-// ScanOneDevice mostly mimics btrfs-progs
-// cmds/rescue-chunk-recover.c:scan_one_device().
-func ScanOneDevice(ctx context.Context, dev *btrfs.Device, sb btrfstree.Superblock) (ScanOneDeviceResult, error) {
-	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-mappings.scan.dev", dev.Name())
-
-	result := ScanOneDeviceResult{
-		FoundNodes: make(map[btrfsvol.LogicalAddr][]btrfsvol.PhysicalAddr),
+func (scanner *deviceScanner) ScanStats() scanStats {
+	return scanStats{
+		NumFoundNodes:       len(scanner.result.FoundNodes),
+		NumFoundChunks:      len(scanner.result.FoundChunks),
+		NumFoundBlockGroups: len(scanner.result.FoundBlockGroups),
+		NumFoundDevExtents:  len(scanner.result.FoundDevExtents),
+		NumFoundExtentCSums: len(scanner.result.FoundExtentCSums),
 	}
+}
 
-	devSize := dev.Size()
-	if sb.NodeSize < sb.SectorSize {
-		return result, fmt.Errorf("node_size(%v) < sector_size(%v)",
-			sb.NodeSize, sb.SectorSize)
-	}
-	if sb.SectorSize != btrfssum.BlockSize {
-		// TODO: probably handle this?
-		return result, fmt.Errorf("sector_size(%v) != btrfssum.BlockSize",
-			sb.SectorSize)
-	}
-	alg := sb.ChecksumType
-	csumSize := alg.Size()
-	numSums := int(devSize / btrfssum.BlockSize)
-	var sums strings.Builder
-	sums.Grow(numSums * csumSize)
+func newDeviceScanner(ctx context.Context, sb btrfstree.Superblock, numBytes btrfsvol.PhysicalAddr, numSectors int) btrfsutil.DeviceScanner[scanStats, ScanOneDeviceResult] {
+	scanner := new(deviceScanner)
+	scanner.alg = sb.ChecksumType
+	scanner.result.FoundNodes = make(map[btrfsvol.LogicalAddr][]btrfsvol.PhysicalAddr)
+	scanner.result.Checksums.ChecksumSize = scanner.alg.Size()
+	scanner.sums.Grow(scanner.result.Checksums.ChecksumSize * numSectors)
+	return scanner
+}
 
-	progressWriter := textui.NewProgress[scanStats](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
-	progress := func(pos btrfsvol.PhysicalAddr) {
-		progressWriter.Set(scanStats{
-			Portion: textui.Portion[btrfsvol.PhysicalAddr]{
-				N: pos,
-				D: devSize,
-			},
-			NumFoundNodes:       len(result.FoundNodes),
-			NumFoundChunks:      len(result.FoundChunks),
-			NumFoundBlockGroups: len(result.FoundBlockGroups),
-			NumFoundDevExtents:  len(result.FoundDevExtents),
-			NumFoundExtentCSums: len(result.FoundExtentCSums),
-		})
+func (scanner *deviceScanner) ScanSector(ctx context.Context, dev *btrfs.Device, paddr btrfsvol.PhysicalAddr) error {
+	sum, err := btrfs.ChecksumPhysical(dev, scanner.alg, paddr)
+	if err != nil {
+		return err
 	}
+	scanner.sums.Write(sum[:scanner.result.Checksums.ChecksumSize])
+	return nil
+}
 
-	var minNextNode btrfsvol.PhysicalAddr
-	for i := 0; i < numSums; i++ {
-		if ctx.Err() != nil {
-			return result, ctx.Err()
-		}
-		pos := btrfsvol.PhysicalAddr(i * btrfssum.BlockSize)
-		progress(pos)
-
-		sum, err := btrfs.ChecksumPhysical(dev, alg, pos)
-		if err != nil {
-			return result, err
-		}
-		sums.Write(sum[:csumSize])
-
-		checkForNode := pos >= minNextNode && pos+btrfsvol.PhysicalAddr(sb.NodeSize) <= devSize
-		if checkForNode {
-			for _, sbAddr := range btrfs.SuperblockAddrs {
-				if sbAddr <= pos && pos < sbAddr+sbSize {
-					checkForNode = false
-					break
-				}
+func (scanner *deviceScanner) ScanNode(ctx context.Context, nodeRef *diskio.Ref[btrfsvol.PhysicalAddr, btrfstree.Node]) error {
+	scanner.result.FoundNodes[nodeRef.Data.Head.Addr] = append(scanner.result.FoundNodes[nodeRef.Data.Head.Addr], nodeRef.Addr)
+	for i, item := range nodeRef.Data.BodyLeaf {
+		switch item.Key.ItemType {
+		case btrfsitem.CHUNK_ITEM_KEY:
+			switch itemBody := item.Body.(type) {
+			case *btrfsitem.Chunk:
+				dlog.Tracef(ctx, "node@%v: item %v: found chunk",
+					nodeRef.Addr, i)
+				scanner.result.FoundChunks = append(scanner.result.FoundChunks, btrfstree.SysChunk{
+					Key:   item.Key,
+					Chunk: *itemBody,
+				})
+			case *btrfsitem.Error:
+				dlog.Errorf(ctx, "node@%v: item %v: error: malformed CHUNK_ITEM: %v",
+					nodeRef.Addr, i, itemBody.Err)
+			default:
+				panic(fmt.Errorf("should not happen: CHUNK_ITEM has unexpected item type: %T", itemBody))
 			}
-		}
-
-		if checkForNode {
-			nodeRef, err := btrfstree.ReadNode[btrfsvol.PhysicalAddr](dev, sb, pos, btrfstree.NodeExpectations{})
-			if err != nil {
-				if !errors.Is(err, btrfstree.ErrNotANode) {
-					dlog.Errorf(ctx, "error: %v", err)
-				}
-			} else {
-				result.FoundNodes[nodeRef.Data.Head.Addr] = append(result.FoundNodes[nodeRef.Data.Head.Addr], nodeRef.Addr)
-				for i, item := range nodeRef.Data.BodyLeaf {
-					switch item.Key.ItemType {
-					case btrfsitem.CHUNK_ITEM_KEY:
-						switch itemBody := item.Body.(type) {
-						case *btrfsitem.Chunk:
-							dlog.Tracef(ctx, "node@%v: item %v: found chunk",
-								nodeRef.Addr, i)
-							result.FoundChunks = append(result.FoundChunks, btrfstree.SysChunk{
-								Key:   item.Key,
-								Chunk: *itemBody,
-							})
-						case *btrfsitem.Error:
-							dlog.Errorf(ctx, "node@%v: item %v: error: malformed CHUNK_ITEM: %v",
-								nodeRef.Addr, i, itemBody.Err)
-						default:
-							panic(fmt.Errorf("should not happen: CHUNK_ITEM has unexpected item type: %T", itemBody))
-						}
-					case btrfsitem.BLOCK_GROUP_ITEM_KEY:
-						switch itemBody := item.Body.(type) {
-						case *btrfsitem.BlockGroup:
-							dlog.Tracef(ctx, "node@%v: item %v: found block group",
-								nodeRef.Addr, i)
-							result.FoundBlockGroups = append(result.FoundBlockGroups, SysBlockGroup{
-								Key: item.Key,
-								BG:  *itemBody,
-							})
-						case *btrfsitem.Error:
-							dlog.Errorf(ctx, "node@%v: item %v: error: malformed BLOCK_GROUP_ITEM: %v",
-								nodeRef.Addr, i, itemBody.Err)
-						default:
-							panic(fmt.Errorf("should not happen: BLOCK_GROUP_ITEM has unexpected item type: %T", itemBody))
-						}
-					case btrfsitem.DEV_EXTENT_KEY:
-						switch itemBody := item.Body.(type) {
-						case *btrfsitem.DevExtent:
-							dlog.Tracef(ctx, "node@%v: item %v: found dev extent",
-								nodeRef.Addr, i)
-							result.FoundDevExtents = append(result.FoundDevExtents, SysDevExtent{
-								Key:    item.Key,
-								DevExt: *itemBody,
-							})
-						case *btrfsitem.Error:
-							dlog.Errorf(ctx, "node@%v: item %v: error: malformed DEV_EXTENT: %v",
-								nodeRef.Addr, i, itemBody.Err)
-						default:
-							panic(fmt.Errorf("should not happen: DEV_EXTENT has unexpected item type: %T", itemBody))
-						}
-					case btrfsitem.EXTENT_CSUM_KEY:
-						switch itemBody := item.Body.(type) {
-						case *btrfsitem.ExtentCSum:
-							dlog.Tracef(ctx, "node@%v: item %v: found csums",
-								nodeRef.Addr, i)
-							result.FoundExtentCSums = append(result.FoundExtentCSums, SysExtentCSum{
-								Generation: nodeRef.Data.Head.Generation,
-								Sums:       *itemBody,
-							})
-						case *btrfsitem.Error:
-							dlog.Errorf(ctx, "node@%v: item %v: error: malformed is EXTENT_CSUM: %v",
-								nodeRef.Addr, i, itemBody.Err)
-						default:
-							panic(fmt.Errorf("should not happen: EXTENT_CSUM has unexpected item type: %T", itemBody))
-						}
-					}
-				}
-				minNextNode = pos + btrfsvol.PhysicalAddr(sb.NodeSize)
+		case btrfsitem.BLOCK_GROUP_ITEM_KEY:
+			switch itemBody := item.Body.(type) {
+			case *btrfsitem.BlockGroup:
+				dlog.Tracef(ctx, "node@%v: item %v: found block group",
+					nodeRef.Addr, i)
+				scanner.result.FoundBlockGroups = append(scanner.result.FoundBlockGroups, SysBlockGroup{
+					Key: item.Key,
+					BG:  *itemBody,
+				})
+			case *btrfsitem.Error:
+				dlog.Errorf(ctx, "node@%v: item %v: error: malformed BLOCK_GROUP_ITEM: %v",
+					nodeRef.Addr, i, itemBody.Err)
+			default:
+				panic(fmt.Errorf("should not happen: BLOCK_GROUP_ITEM has unexpected item type: %T", itemBody))
+			}
+		case btrfsitem.DEV_EXTENT_KEY:
+			switch itemBody := item.Body.(type) {
+			case *btrfsitem.DevExtent:
+				dlog.Tracef(ctx, "node@%v: item %v: found dev extent",
+					nodeRef.Addr, i)
+				scanner.result.FoundDevExtents = append(scanner.result.FoundDevExtents, SysDevExtent{
+					Key:    item.Key,
+					DevExt: *itemBody,
+				})
+			case *btrfsitem.Error:
+				dlog.Errorf(ctx, "node@%v: item %v: error: malformed DEV_EXTENT: %v",
+					nodeRef.Addr, i, itemBody.Err)
+			default:
+				panic(fmt.Errorf("should not happen: DEV_EXTENT has unexpected item type: %T", itemBody))
+			}
+		case btrfsitem.EXTENT_CSUM_KEY:
+			switch itemBody := item.Body.(type) {
+			case *btrfsitem.ExtentCSum:
+				dlog.Tracef(ctx, "node@%v: item %v: found csums",
+					nodeRef.Addr, i)
+				scanner.result.FoundExtentCSums = append(scanner.result.FoundExtentCSums, SysExtentCSum{
+					Generation: nodeRef.Data.Head.Generation,
+					Sums:       *itemBody,
+				})
+			case *btrfsitem.Error:
+				dlog.Errorf(ctx, "node@%v: item %v: error: malformed is EXTENT_CSUM: %v",
+					nodeRef.Addr, i, itemBody.Err)
+			default:
+				panic(fmt.Errorf("should not happen: EXTENT_CSUM has unexpected item type: %T", itemBody))
 			}
-			btrfstree.FreeNodeRef(nodeRef)
 		}
 	}
-	progress(devSize)
-	progressWriter.Done()
-
-	result.Checksums = btrfssum.SumRun[btrfsvol.PhysicalAddr]{
-		ChecksumSize: csumSize,
-		Sums:         btrfssum.ShortSum(sums.String()),
-	}
+	return nil
+}
 
-	return result, nil
+func (scanner *deviceScanner) ScanDone(ctx context.Context) (ScanOneDeviceResult, error) {
+	scanner.result.Checksums.Sums = btrfssum.ShortSum(scanner.sums.String())
+	return scanner.result, nil
 }
diff --git a/lib/btrfsutil/scan.go b/lib/btrfsutil/scan.go
new file mode 100644
index 0000000..97220aa
--- /dev/null
+++ b/lib/btrfsutil/scan.go
@@ -0,0 +1,148 @@
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package btrfsutil
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/datawire/dlib/dgroup"
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfssum"
+	"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"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+var sbSize = btrfsvol.PhysicalAddr(binstruct.StaticSize(btrfstree.Superblock{}))
+
+type DeviceScannerFactory[Stats comparable, Result any] func(ctx context.Context, sb btrfstree.Superblock, numBytes btrfsvol.PhysicalAddr, numSectors int) DeviceScanner[Stats, Result]
+
+type DeviceScanner[Stats comparable, Result any] interface {
+	ScanStats() Stats
+	ScanSector(ctx context.Context, dev *btrfs.Device, paddr btrfsvol.PhysicalAddr) error
+	ScanNode(ctx context.Context, nodeRef *diskio.Ref[btrfsvol.PhysicalAddr, btrfstree.Node]) error
+	ScanDone(ctx context.Context) (Result, error)
+}
+
+type scanStats[T comparable] struct {
+	portion textui.Portion[btrfsvol.PhysicalAddr]
+	stats   T
+}
+
+func (s scanStats[T]) String() string {
+	return textui.Sprintf("scanned %v (%v)",
+		s.portion, s.stats)
+}
+
+func ScanDevices[Stats comparable, Result any](ctx context.Context, fs *btrfs.FS, newScanner DeviceScannerFactory[Stats, Result]) (map[btrfsvol.DeviceID]Result, error) {
+	grp := dgroup.NewGroup(ctx, dgroup.GroupConfig{})
+	var mu sync.Mutex
+	result := make(map[btrfsvol.DeviceID]Result)
+	for id, dev := range fs.LV.PhysicalVolumes() {
+		id := id
+		dev := dev
+		grp.Go(fmt.Sprintf("dev-%d", id), func(ctx context.Context) error {
+			devResult, err := ScanOneDevice[Stats, Result](ctx, dev, newScanner)
+			if err != nil {
+				return err
+			}
+			mu.Lock()
+			result[id] = devResult
+			mu.Unlock()
+			return nil
+		})
+	}
+	if err := grp.Wait(); err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
+func ScanOneDevice[Stats comparable, Result any](ctx context.Context, dev *btrfs.Device, newScanner DeviceScannerFactory[Stats, Result]) (Result, error) {
+	ctx = dlog.WithField(ctx, "scandevices.dev", dev.Name())
+
+	sb, err := dev.Superblock()
+	if err != nil {
+		var zero Result
+		return zero, err
+	}
+	numBytes := dev.Size()
+	if sb.NodeSize < sb.SectorSize {
+		var zero Result
+		return zero, fmt.Errorf("node_size(%v) < sector_size(%v)",
+			sb.NodeSize, sb.SectorSize)
+	}
+	if sb.SectorSize != btrfssum.BlockSize {
+		// TODO: probably handle this?
+		var zero Result
+		return zero, fmt.Errorf("sector_size(%v) != btrfssum.BlockSize",
+			sb.SectorSize)
+	}
+	numSectors := int(numBytes / btrfssum.BlockSize)
+
+	scanner := newScanner(ctx, *sb, numBytes, numSectors)
+
+	progressWriter := textui.NewProgress[scanStats[Stats]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+	var stats scanStats[Stats]
+	stats.portion.D = numBytes
+
+	var minNextNode btrfsvol.PhysicalAddr
+	for i := 0; i < numSectors; i++ {
+		if ctx.Err() != nil {
+			var zero Result
+			return zero, ctx.Err()
+		}
+		pos := btrfsvol.PhysicalAddr(i * btrfssum.BlockSize)
+		stats.portion.N = pos
+		stats.stats = scanner.ScanStats()
+		progressWriter.Set(stats)
+
+		if err := scanner.ScanSector(ctx, dev, pos); err != nil {
+			var zero Result
+			return zero, err
+		}
+
+		checkForNode := pos >= minNextNode && pos+btrfsvol.PhysicalAddr(sb.NodeSize) <= numBytes
+		if checkForNode {
+			for _, sbAddr := range btrfs.SuperblockAddrs {
+				if sbAddr <= pos && pos < sbAddr+sbSize {
+					checkForNode = false
+					break
+				}
+			}
+		}
+
+		if checkForNode {
+			nodeRef, err := btrfstree.ReadNode[btrfsvol.PhysicalAddr](dev, *sb, pos, btrfstree.NodeExpectations{})
+			if err != nil {
+				if !errors.Is(err, btrfstree.ErrNotANode) {
+					dlog.Errorf(ctx, "error: %v", err)
+				}
+			} else {
+				if err := scanner.ScanNode(ctx, nodeRef); err != nil {
+					var zero Result
+					return zero, err
+				}
+				minNextNode = pos + btrfsvol.PhysicalAddr(sb.NodeSize)
+			}
+			btrfstree.FreeNodeRef(nodeRef)
+		}
+	}
+
+	stats.portion.N = numBytes
+	stats.stats = scanner.ScanStats()
+	progressWriter.Set(stats)
+	progressWriter.Done()
+
+	return scanner.ScanDone(ctx)
+}