Exit early from eachDiff if the context was cancelledfix/exit-early-from-each-diff-if-ctx-is-cancelled

2 files changed, 10 insertions, 0 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c580e7acb..910fbd681 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@ UNRELEASED
 
 - Fix incorrect parsing of diff chunks starting with ++ or --
   https://gitlab.com/gitlab-org/gitaly/merge_requests/385
+- Exit early from eachDiff if the context was cancelled
+  https://gitlab.com/gitlab-org/gitaly/merge_requests/390
 - Implement Raw{Diff,Patch} RPCs
   https://gitlab.com/gitlab-org/gitaly/merge_requests/381
 
diff --git a/internal/service/diff/commit.go b/internal/service/diff/commit.go
index 59b478809..81e6a96d5 100644
--- a/internal/service/diff/commit.go
+++ b/internal/service/diff/commit.go
@@ -223,6 +223,14 @@ func eachDiff(ctx context.Context, rpc string, repo *pb.Repository, cmdArgs []st
 		return grpc.Errorf(codes.Internal, "%s: parse failure: %v", rpc, err)
 	}
 
+	// The context could be cancelled early because the client ended the request
+	// (e.g. it hit a limit), so we exit early here because the command has already
+	// been terminated as a cancellation result and `Wait` below would return an error
+	// we don't want.
+	if ctx.Err() == context.Canceled {
+		return nil
+	}
+
 	if err := cmd.Wait(); err != nil {
 		return grpc.Errorf(codes.Unavailable, "%s: %v", rpc, err)
 	}