From 8efc82d0b1a167830970135c78d173667080b116 Mon Sep 17 00:00:00 2001 From: Luke Shumaker Date: Wed, 28 Dec 2022 18:49:09 -0700 Subject: rebuildnodes: Support graceful shutdown --- .golangci.yml | 1 - cmd/btrfs-rec/inspect_rebuildnodes.go | 18 +++++-- cmd/btrfs-rec/util.go | 17 ++++++- .../rebuildnodes/btrees/rebuilt_btrees.go | 9 ++-- .../btrfsinspect/rebuildnodes/rebuild.go | 58 +++++++++++++--------- lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go | 18 ++++--- 6 files changed, 81 insertions(+), 40 deletions(-) diff --git a/.golangci.yml b/.golangci.yml index 6b9a830..355664f 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -51,7 +51,6 @@ linters: - revive - testpackage - thelper - - unparam - varnamelen - wrapcheck linters-settings: diff --git a/cmd/btrfs-rec/inspect_rebuildnodes.go b/cmd/btrfs-rec/inspect_rebuildnodes.go index e61e6d2..9c86c3a 100644 --- a/cmd/btrfs-rec/inspect_rebuildnodes.go +++ b/cmd/btrfs-rec/inspect_rebuildnodes.go @@ -33,21 +33,31 @@ func init() { } dlog.Infof(ctx, "... done reading %q", args[0]) - rebuiltNodes, err := rebuildnodes.RebuildNodes(ctx, fs, nodeScanResults) + rebuilder, err := rebuildnodes.NewRebuilder(ctx, fs, nodeScanResults) if err != nil { return err } - dlog.Info(ctx, "Writing re-built nodes to stdout...") - if err := writeJSONFile(os.Stdout, rebuiltNodes, lowmemjson.ReEncoder{ + dlog.Info(ctx, "Rebuilding node tree...") + rebuildErr := rebuilder.Rebuild(ctx) + dst := os.Stdout + if rebuildErr != nil { + dst = os.Stderr + dlog.Errorf(ctx, "rebuild error: %v", rebuildErr) + } + dlog.Infof(ctx, "Writing re-built nodes to %s...", dst.Name()) + if err := writeJSONFile(dst, rebuilder.ListRoots(), lowmemjson.ReEncoder{ Indent: "\t", ForceTrailingNewlines: true, }); err != nil { + if rebuildErr != nil { + return rebuildErr + } return err } dlog.Info(ctx, "... done writing") - return nil + return rebuildErr }, }) } diff --git a/cmd/btrfs-rec/util.go b/cmd/btrfs-rec/util.go index ffc03cc..9a0d60c 100644 --- a/cmd/btrfs-rec/util.go +++ b/cmd/btrfs-rec/util.go @@ -18,6 +18,7 @@ import ( ) type runeScanner struct { + ctx context.Context //nolint:containedctx // For detecting shutdown from methods progress textui.Portion[int64] progressWriter *textui.Progress[textui.Portion[int64]] unreadCnt uint64 @@ -31,6 +32,7 @@ func newRuneScanner(ctx context.Context, fh *os.File) (*runeScanner, error) { return nil, err } ret := &runeScanner{ + ctx: ctx, progress: textui.Portion[int64]{ D: fi.Size(), }, @@ -42,6 +44,9 @@ func newRuneScanner(ctx context.Context, fh *os.File) (*runeScanner, error) { } func (rs *runeScanner) ReadRune() (r rune, size int, err error) { + if err := rs.ctx.Err(); err != nil { + return 0, 0, err + } r, size, err = rs.reader.ReadRune() if rs.unreadCnt > 0 { rs.unreadCnt-- @@ -53,8 +58,14 @@ func (rs *runeScanner) ReadRune() (r rune, size int, err error) { } func (rs *runeScanner) UnreadRune() error { + if err := rs.ctx.Err(); err != nil { + return err + } + if err := rs.reader.UnreadRune(); err != nil { + return err + } rs.unreadCnt++ - return rs.reader.UnreadRune() + return nil } func (rs *runeScanner) Close() error { @@ -69,6 +80,9 @@ func readJSONFile[T any](ctx context.Context, filename string) (T, error) { return zero, err } buf, err := newRuneScanner(dlog.WithField(ctx, "btrfs.read-json-file", filename), fh) + defer func() { + _ = buf.Close() + }() if err != nil { var zero T return zero, err @@ -78,7 +92,6 @@ func readJSONFile[T any](ctx context.Context, filename string) (T, error) { var zero T return zero, err } - _ = buf.Close() return ret, nil } diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go index b1ae7be..ffe225a 100644 --- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go +++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/rebuilt_btrees.go @@ -305,7 +305,7 @@ func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) { progress() for _, node := range maps.SortedKeys(graph.Nodes) { - tree.indexNode(graph, node, nodeToRoots, progress, nil) + tree.indexNode(ctx, graph, node, nodeToRoots, progress, nil) } progressWriter.Done() @@ -317,8 +317,11 @@ func (tree *rebuiltTree) indexLeafs(ctx context.Context, graph pkggraph.Graph) { } } -func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) { +func (tree *rebuiltTree) indexNode(ctx context.Context, graph pkggraph.Graph, node btrfsvol.LogicalAddr, index map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr], progress func(), stack []btrfsvol.LogicalAddr) { defer progress() + if err := ctx.Err(); err != nil { + return + } if _, done := index[node]; done { return } @@ -339,7 +342,7 @@ func (tree *rebuiltTree) indexNode(graph pkggraph.Graph, node btrfsvol.LogicalAd return !tree.isOwnerOK(graph.Nodes[kp.FromNode].Owner, graph.Nodes[kp.FromNode].Generation) }) for _, kp := range kps { - tree.indexNode(graph, kp.FromNode, index, progress, stack) + tree.indexNode(ctx, graph, kp.FromNode, index, progress, stack) if len(index[kp.FromNode]) > 0 { if roots == nil { roots = make(containers.Set[btrfsvol.LogicalAddr]) diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go index b6c1359..26f2a44 100644 --- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go +++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go @@ -44,47 +44,38 @@ func (o keyAndTree) String() string { } type rebuilder struct { - sb btrfstree.Superblock - rebuilt *btrees.RebuiltTrees - + sb btrfstree.Superblock graph graph.Graph keyIO *keyio.Handle + rebuilt *btrees.RebuiltTrees + curKey keyAndTree treeQueue containers.Set[btrfsprim.ObjID] itemQueue containers.Set[keyAndTree] augmentQueue map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int } -func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr], error) { - _ctx := ctx +type Rebuilder interface { + Rebuild(context.Context) error + ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] +} - ctx = dlog.WithField(_ctx, "btrfsinspect.rebuild-nodes.step", "read-fs-data") - dlog.Info(ctx, "Reading superblock...") - sb, err := fs.Superblock() - if err != nil { - return nil, err - } - nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging +func NewRebuilder(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (Rebuilder, error) { + ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.step", "read-fs-data") + sb, nodeGraph, keyIO, err := ScanDevices(ctx, fs, nodeScanResults) // ScanDevices does its own logging if err != nil { return nil, err } - ctx = dlog.WithField(_ctx, "btrfsinspect.rebuild-nodes.step", "rebuild") - dlog.Info(ctx, "Rebuilding node tree...") o := &rebuilder{ - sb: *sb, - + sb: sb, graph: nodeGraph, keyIO: keyIO, } - o.rebuilt = btrees.NewRebuiltTrees(*sb, nodeGraph, keyIO, + o.rebuilt = btrees.NewRebuiltTrees(sb, nodeGraph, keyIO, o.cbAddedItem, o.cbLookupRoot, o.cbLookupUUID) - if err := o.rebuild(ctx); err != nil { - return nil, err - } - - return o.rebuilt.ListRoots(), nil + return o, nil } func (o *rebuilder) ioErr(ctx context.Context, err error) { @@ -93,7 +84,13 @@ func (o *rebuilder) ioErr(ctx context.Context, err error) { panic(err) } -func (o *rebuilder) rebuild(_ctx context.Context) error { +func (o *rebuilder) ListRoots() map[btrfsprim.ObjID]containers.Set[btrfsvol.LogicalAddr] { + return o.rebuilt.ListRoots() +} + +func (o *rebuilder) Rebuild(_ctx context.Context) error { + _ctx = dlog.WithField(_ctx, "btrfsinspect.rebuild-nodes.step", "rebuild") + // Initialize o.itemQueue = make(containers.Set[keyAndTree]) o.augmentQueue = make(map[btrfsprim.ObjID][]map[btrfsvol.LogicalAddr]int) @@ -116,6 +113,9 @@ func (o *rebuilder) rebuild(_ctx context.Context) error { // Because trees can be wildly different sizes, it's impossible to have a meaningful // progress percentage here. for _, treeID := range maps.SortedKeys(treeQueue) { + if err := _ctx.Err(); err != nil { + return err + } o.rebuilt.AddTree(stepCtx, treeID) } @@ -134,6 +134,10 @@ func (o *rebuilder) rebuild(_ctx context.Context) error { itemCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.process.item", key) progress.N = i progressWriter.Set(progress) + if err := _ctx.Err(); err != nil { + progressWriter.Done() + return err + } o.curKey = key itemBody, ok := o.rebuilt.Load(itemCtx, key.TreeID, key.Key) if !ok { @@ -157,6 +161,9 @@ func (o *rebuilder) rebuild(_ctx context.Context) error { progress.N = 0 progress.D = 0 for _, treeID := range maps.SortedKeys(o.augmentQueue) { + if err := _ctx.Err(); err != nil { + return err + } treeCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.augment.tree", treeID) resolvedAugments[treeID] = o.resolveTreeAugments(treeCtx, o.augmentQueue[treeID]) progress.D += len(resolvedAugments[treeID]) @@ -167,6 +174,11 @@ func (o *rebuilder) rebuild(_ctx context.Context) error { for _, treeID := range maps.SortedKeys(resolvedAugments) { treeCtx := dlog.WithField(stepCtx, "btrfsinspect.rebuild-nodes.rebuild.augment.tree", treeID) for _, nodeAddr := range maps.SortedKeys(resolvedAugments[treeID]) { + if err := _ctx.Err(); err != nil { + progressWriter.Set(progress) + progressWriter.Done() + return err + } progressWriter.Set(progress) o.rebuilt.AddRoot(treeCtx, treeID, nodeAddr) progress.N++ diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go index 7e96e29..7e19802 100644 --- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go +++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go @@ -20,14 +20,15 @@ import ( "git.lukeshu.com/btrfs-progs-ng/lib/textui" ) -func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (graph.Graph, *keyio.Handle, error) { - dlog.Infof(ctx, "Reading node data from FS...") - +func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (btrfstree.Superblock, graph.Graph, *keyio.Handle, error) { + dlog.Info(ctx, "Reading superblock...") sb, err := fs.Superblock() if err != nil { - return graph.Graph{}, nil, err + return btrfstree.Superblock{}, graph.Graph{}, nil, err } + dlog.Infof(ctx, "Reading node data from FS...") + var stats textui.Portion[int] stats.D = countNodes(scanResults) progressWriter := textui.NewProgress[textui.Portion[int]]( @@ -40,11 +41,14 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca progressWriter.Set(stats) for _, devResults := range scanResults { for laddr := range devResults.FoundNodes { + if err := ctx.Err(); err != nil { + return btrfstree.Superblock{}, graph.Graph{}, nil, err + } nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{ LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr}, }) if err != nil { - return graph.Graph{}, nil, err + return btrfstree.Superblock{}, graph.Graph{}, nil, err } nodeGraph.InsertNode(nodeRef) @@ -61,9 +65,9 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca dlog.Info(ctx, "... done reading node data") if err := nodeGraph.FinalCheck(ctx, fs, *sb); err != nil { - return graph.Graph{}, nil, err + return btrfstree.Superblock{}, graph.Graph{}, nil, err } keyIO.SetGraph(*nodeGraph) - return *nodeGraph, keyIO, nil + return *sb, *nodeGraph, keyIO, nil } -- cgit v1.2.3-54-g00ecf