Improve command shutdown on context cancellation

cmd/restic/cmd_check.go

@@ -219,6 +219,9 @@ func runCheck(ctx context.Context, opts CheckOptions, gopts GlobalOptions, args
 	Verbosef("load indexes\n")
 	bar := newIndexProgress(gopts.Quiet, gopts.JSON)
 	hints, errs := chkr.LoadIndex(ctx, bar)
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	errorsFound := false
 	suggestIndexRebuild := false
@@ -280,6 +283,9 @@ func runCheck(ctx context.Context, opts CheckOptions, gopts GlobalOptions, args
 	if orphanedPacks > 0 {
 		Verbosef("%d additional files were found in the repo, which likely contain duplicate data.\nThis is non-critical, you can run `restic prune` to correct this.\n", orphanedPacks)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	Verbosef("check snapshots, trees and blobs\n")
 	errChan = make(chan error)
@@ -313,6 +319,9 @@ func runCheck(ctx context.Context, opts CheckOptions, gopts GlobalOptions, args
 	// Must happen after `errChan` is read from in the above loop to avoid
 	// deadlocking in the case of errors.
 	wg.Wait()
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	if opts.CheckUnused {
 		for _, id := range chkr.UnusedBlobs(ctx) {
@@ -392,10 +401,13 @@ func runCheck(ctx context.Context, opts CheckOptions, gopts GlobalOptions, args
 		doReadData(packs)
 	}
 
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
 	if errorsFound {
 		return errors.Fatal("repository contains errors")
 	}
-
 	Verbosef("no errors were found\n")
 
 	return nil

cmd/restic/cmd_copy.go

@@ -103,6 +103,9 @@ func runCopy(ctx context.Context, opts CopyOptions, gopts GlobalOptions, args []
 		// also consider identical snapshot copies
 		dstSnapshotByOriginal[*sn.ID()] = append(dstSnapshotByOriginal[*sn.ID()], sn)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	// remember already processed trees across all snapshots
 	visitedTrees := restic.NewIDSet()
@@ -147,7 +150,7 @@ func runCopy(ctx context.Context, opts CopyOptions, gopts GlobalOptions, args []
 		}
 		Verbosef("snapshot %s saved\n", newID.Str())
 	}
-	return nil
+	return ctx.Err()
 }
 
 func similarSnapshots(sna *restic.Snapshot, snb *restic.Snapshot) bool {

cmd/restic/cmd_find.go

@@ -608,6 +608,9 @@ func runFind(ctx context.Context, opts FindOptions, gopts GlobalOptions, args []
 	for sn := range FindFilteredSnapshots(ctx, snapshotLister, repo, &opts.SnapshotFilter, opts.Snapshots) {
 		filteredSnapshots = append(filteredSnapshots, sn)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	sort.Slice(filteredSnapshots, func(i, j int) bool {
 		return filteredSnapshots[i].Time.Before(filteredSnapshots[j].Time)

cmd/restic/cmd_forget.go

@@ -188,6 +188,9 @@ func runForget(ctx context.Context, opts ForgetOptions, pruneOptions PruneOption
 	for sn := range FindFilteredSnapshots(ctx, repo, repo, &opts.SnapshotFilter, args) {
 		snapshots = append(snapshots, sn)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	var jsonGroups []*ForgetGroup
 
@@ -270,6 +273,10 @@ func runForget(ctx context.Context, opts ForgetOptions, pruneOptions PruneOption
 		}
 	}
 
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
 	if len(removeSnIDs) > 0 {
 		if !opts.DryRun {
 			bar := printer.NewCounter("files deleted")

cmd/restic/cmd_prune.go

@@ -197,6 +197,9 @@ func runPruneWithRepo(ctx context.Context, opts PruneOptions, gopts GlobalOption
 	if err != nil {
 		return err
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	if popts.DryRun {
 		printer.P("\nWould have made the following changes:")

cmd/restic/cmd_recover.go

@@ -66,11 +66,17 @@ func runRecover(ctx context.Context, gopts GlobalOptions) error {
 			trees[blob.Blob.ID] = false
 		}
 	})
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	Verbosef("load %d trees\n", len(trees))
 	bar = newProgressMax(!gopts.Quiet, uint64(len(trees)), "trees loaded")
 	for id := range trees {
 		tree, err := restic.LoadTree(ctx, repo, id)
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
 		if err != nil {
 			Warnf("unable to load tree %v: %v\n", id.Str(), err)
 			continue

cmd/restic/cmd_repair_snapshots.go

@@ -145,6 +145,9 @@ func runRepairSnapshots(ctx context.Context, gopts GlobalOptions, opts RepairOpt
 			changedCount++
 		}
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	Verbosef("\n")
 	if changedCount == 0 {

cmd/restic/cmd_rewrite.go

@@ -294,6 +294,9 @@ func runRewrite(ctx context.Context, opts RewriteOptions, gopts GlobalOptions, a
 			changedCount++
 		}
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	Verbosef("\n")
 	if changedCount == 0 {

cmd/restic/cmd_snapshots.go

@@ -69,6 +69,9 @@ func runSnapshots(ctx context.Context, opts SnapshotOptions, gopts GlobalOptions
 	for sn := range FindFilteredSnapshots(ctx, repo, repo, &opts.SnapshotFilter, args) {
 		snapshots = append(snapshots, sn)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 	snapshotGroups, grouped, err := restic.GroupSnapshots(snapshots, opts.GroupBy)
 	if err != nil {
 		return err

cmd/restic/cmd_stats.go

@@ -117,9 +117,8 @@ func runStats(ctx context.Context, opts StatsOptions, gopts GlobalOptions, args
 			return fmt.Errorf("error walking snapshot: %v", err)
 		}
 	}
-
-	if err != nil {
-		return err
+	if ctx.Err() != nil {
+		return ctx.Err()
 	}
 
 	if opts.countMode == countModeRawData {

cmd/restic/cmd_tag.go

@@ -122,6 +122,9 @@ func runTag(ctx context.Context, opts TagOptions, gopts GlobalOptions, args []st
 			changeCnt++
 		}
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 	if changeCnt == 0 {
 		Verbosef("no snapshots were modified\n")
 	} else {