Add more checks for canceled contexts

cmd/restic/cmd_diff.go

@@ -177,6 +177,10 @@ func (c *Comparer) printDir(ctx context.Context, mode string, stats *DiffStat, b
 	}
 
 	for _, node := range tree.Nodes {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		name := path.Join(prefix, node.Name)
 		if node.Type == "dir" {
 			name += "/"
@@ -204,6 +208,10 @@ func (c *Comparer) collectDir(ctx context.Context, blobs restic.BlobSet, id rest
 	}
 
 	for _, node := range tree.Nodes {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		addBlobs(blobs, node)
 
 		if node.Type == "dir" {
@@ -255,6 +263,10 @@ func (c *Comparer) diffTree(ctx context.Context, stats *DiffStatsContainer, pref
 	tree1Nodes, tree2Nodes, names := uniqueNodeNames(tree1, tree2)
 
 	for _, name := range names {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		node1, t1 := tree1Nodes[name]
 		node2, t2 := tree2Nodes[name]
 

cmd/restic/cmd_dump.go

@@ -85,6 +85,10 @@ func printFromTree(ctx context.Context, tree *restic.Tree, repo restic.BlobLoade
 	item := filepath.Join(prefix, pathComponents[0])
 	l := len(pathComponents)
 	for _, node := range tree.Nodes {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		// If dumping something in the highest level it will just take the
 		// first item it finds and dump that according to the switch case below.
 		if node.Name == pathComponents[0] {

cmd/restic/cmd_find.go

@@ -377,6 +377,10 @@ func (f *Finder) findIDs(ctx context.Context, sn *restic.Snapshot) error {
 
 		if node.Type == "file" && f.blobIDs != nil {
 			for _, id := range node.Content {
+				if ctx.Err() != nil {
+					return ctx.Err()
+				}
+
 				idStr := id.String()
 				if _, ok := f.blobIDs[idStr]; !ok {
 					// Look for short ID form

cmd/restic/cmd_forget.go

@@ -246,6 +246,10 @@ func runForget(ctx context.Context, opts ForgetOptions, pruneOptions PruneOption
 		printer.P("Applying Policy: %v\n", policy)
 
 		for k, snapshotGroup := range snapshotGroups {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+
 			if gopts.Verbose >= 1 && !gopts.JSON {
 				err = PrintSnapshotGroupHeader(globalOptions.stdout, k)
 				if err != nil {

cmd/restic/cmd_recover.go

@@ -118,6 +118,10 @@ func runRecover(ctx context.Context, gopts GlobalOptions) error {
 		return nil
 	}
 
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
 	tree := restic.NewTree(len(roots))
 	for id := range roots {
 		var subtreeID = id

cmd/restic/cmd_snapshots.go

@@ -81,6 +81,10 @@ func runSnapshots(ctx context.Context, opts SnapshotOptions, gopts GlobalOptions
 	}
 
 	for k, list := range snapshotGroups {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		if opts.Last {
 			// This branch should be removed in the same time
 			// that --last.
@@ -101,6 +105,10 @@ func runSnapshots(ctx context.Context, opts SnapshotOptions, gopts GlobalOptions
 	}
 
 	for k, list := range snapshotGroups {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
 		if grouped {
 			err := PrintSnapshotGroupHeader(globalOptions.stdout, k)
 			if err != nil {