Improve command shutdown on context cancellation

internal/archiver/archiver.go

@@ -380,6 +380,7 @@ func (fn *FutureNode) take(ctx context.Context) futureNodeResult {
 			return res
 		}
 	case <-ctx.Done():
+		return futureNodeResult{err: ctx.Err()}
 	}
 	return futureNodeResult{err: errors.Errorf("no result")}
 }

internal/archiver/tree_saver.go

@@ -90,6 +90,10 @@ func (s *TreeSaver) save(ctx context.Context, job *saveTreeJob) (*restic.Node, I
 		// return the error if it wasn't ignored
 		if fnr.err != nil {
 			debug.Log("err for %v: %v", fnr.snPath, fnr.err)
+			if fnr.err == context.Canceled {
+				return nil, stats, fnr.err
+			}
+
 			fnr.err = s.errFn(fnr.target, fnr.err)
 			if fnr.err == nil {
 				// ignore error

internal/index/master_index.go

@@ -320,6 +320,9 @@ func (mi *MasterIndex) Save(ctx context.Context, repo restic.Repository, exclude
 					newIndex = NewIndex()
 				}
 			}
+			if wgCtx.Err() != nil {
+				return wgCtx.Err()
+			}
 		}
 
 		err := newIndex.AddToSupersedes(extraObsolete...)

internal/repository/prune.go

@@ -130,6 +130,9 @@ func PlanPrune(ctx context.Context, opts PruneOptions, repo restic.Repository, g
 			}
 			keepBlobs.Delete(blob.BlobHandle)
 		})
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 
 		if keepBlobs.Len() < blobCount/2 {
 			// replace with copy to shrink map to necessary size if there's a chance to benefit
@@ -166,6 +169,9 @@ func packInfoFromIndex(ctx context.Context, idx restic.MasterIndex, usedBlobs re
 			usedBlobs[bh] = count
 		}
 	})
+	if ctx.Err() != nil {
+		return nil, nil, ctx.Err()
+	}
 
 	// Check if all used blobs have been found in index
 	missingBlobs := restic.NewBlobSet()
@@ -240,6 +246,9 @@ func packInfoFromIndex(ctx context.Context, idx restic.MasterIndex, usedBlobs re
 		// update indexPack
 		indexPack[blob.PackID] = ip
 	})
+	if ctx.Err() != nil {
+		return nil, nil, ctx.Err()
+	}
 
 	// if duplicate blobs exist, those will be set to either "used" or "unused":
 	// - mark only one occurrence of duplicate blobs as used
@@ -286,6 +295,9 @@ func packInfoFromIndex(ctx context.Context, idx restic.MasterIndex, usedBlobs re
 			indexPack[blob.PackID] = ip
 		})
 	}
+	if ctx.Err() != nil {
+		return nil, nil, ctx.Err()
+	}
 
 	// Sanity check. If no duplicates exist, all blobs have value 1. After handling
 	// duplicates, this also applies to duplicates.
@@ -528,6 +540,9 @@ func (plan *PrunePlan) Execute(ctx context.Context, printer progress.Printer) (e
 		printer.P("deleting unreferenced packs\n")
 		_ = deleteFiles(ctx, true, repo, plan.removePacksFirst, restic.PackFile, printer)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	if len(plan.repackPacks) != 0 {
 		printer.P("repacking packs\n")
@@ -578,6 +593,9 @@ func (plan *PrunePlan) Execute(ctx context.Context, printer progress.Printer) (e
 		printer.P("removing %d old packs\n", len(plan.removePacks))
 		_ = deleteFiles(ctx, true, repo, plan.removePacks, restic.PackFile, printer)
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	if plan.opts.UnsafeRecovery {
 		err = rebuildIndexFiles(ctx, repo, plan.ignorePacks, nil, true, printer)

internal/repository/repack.go

@@ -72,7 +72,7 @@ func repack(ctx context.Context, repo restic.Repository, dstRepo restic.Reposito
 				return wgCtx.Err()
 			}
 		}
-		return nil
+		return wgCtx.Err()
 	})
 
 	worker := func() error {

internal/repository/repository.go

@@ -713,6 +713,9 @@ func (r *Repository) LoadIndex(ctx context.Context, p *progress.Counter) error {
 			return errors.New("index uses feature not supported by repository version 1")
 		}
 	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 
 	// remove index files from the cache which have been removed in the repo
 	return r.prepareCache()