From e10082d2a7e24c66fbe83eb94f0b532882141698 Mon Sep 17 00:00:00 2001
From: Chris Cotter <cjcotter@google.com>
Date: Wed, 5 Jan 2022 20:23:46 -0500
Subject: [PATCH] fix(internal/gensupport): check ctx in chunk retry (#1364)

The select statement is non-deterministic, so currently, if the
pause is completed and ALSO the context has been canceled or
timeout elapsed, a request may still occur. This PR prevents that
circumstance from occurring.

Also removed something in a test that seems to be a workaround for
the race condition.

Inspired by #1359 & #1358.
---
 internal/gensupport/resumable.go      | 16 ++++++++++++++++
 internal/gensupport/resumable_test.go |  3 ---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/internal/gensupport/resumable.go b/internal/gensupport/resumable.go
index 0fb74606c3a..6796a07c984 100644
--- a/internal/gensupport/resumable.go
+++ b/internal/gensupport/resumable.go
@@ -177,6 +177,22 @@ func (rx *ResumableUpload) Upload(ctx context.Context) (resp *http.Response, err
 				return prepareReturn(resp, err)
 			}
 
+			// Check for context cancellation or timeout once more. If more than one
+			// case in the select statement above was satisfied at the same time, Go
+			// will choose one arbitrarily.
+			// That can cause an operation to go through even if the context was
+			// canceled before or the timeout was reached.
+			select {
+			case <-ctx.Done():
+				if err == nil {
+					err = ctx.Err()
+				}
+				return prepareReturn(resp, err)
+			case <-quitAfter:
+				return prepareReturn(resp, err)
+			default:
+			}
+
 			resp, err = rx.transferChunk(ctx)
 
 			var status int
diff --git a/internal/gensupport/resumable_test.go b/internal/gensupport/resumable_test.go
index 44544f7d71c..bffd88098e2 100644
--- a/internal/gensupport/resumable_test.go
+++ b/internal/gensupport/resumable_test.go
@@ -213,9 +213,6 @@ func TestCancelUploadFast(t *testing.T) {
 
 	tr := &interruptibleTransport{
 		buf: make([]byte, 0, mediaSize),
-		// Shouldn't really need an event, but sometimes the test loses the
-		// race. So, this is just a filler event.
-		events: []event{{"bytes 0-9/*", http.StatusServiceUnavailable}},
 	}
 
 	pr := progressRecorder{}