From 2eedfd9dda7e44d87b73f6799b0e9488e75d489d Mon Sep 17 00:00:00 2001
From: akartasov <agneum@gmail.com>
Date: Fri, 2 Oct 2020 15:53:33 +0700
Subject: [PATCH] fix(internal/gensupport): allow user to provide his own
 retryDeadline for uploading (#685)

---
 internal/gensupport/resumable.go      | 16 +++++++++++++---
 internal/gensupport/resumable_test.go |  6 +-----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/internal/gensupport/resumable.go b/internal/gensupport/resumable.go
index edc87ec24f6..3d3f9ba5ba2 100644
--- a/internal/gensupport/resumable.go
+++ b/internal/gensupport/resumable.go
@@ -24,8 +24,7 @@ type Backoff interface {
 
 // These are declared as global variables so that tests can overwrite them.
 var (
-	retryDeadline = 32 * time.Second
-	backoff       = func() Backoff {
+	backoff = func() Backoff {
 		return &gax.Backoff{Initial: 100 * time.Millisecond}
 	}
 	// isRetryable is a platform-specific hook, specified in retryable_linux.go
@@ -38,6 +37,12 @@ const (
 	// should be retried.
 	// https://cloud.google.com/storage/docs/json_api/v1/status-codes#standardcodes
 	statusTooManyRequests = 429
+
+	// retryDeadlineKey defines a context key to override the default retry deadline.
+	retryDeadlineKey = "retryDeadline"
+
+	// defaultRetryDeadline defines the default deadline duration for a single chunk uploading.
+	defaultRetryDeadline = 32 * time.Second
 )
 
 // ResumableUpload is used by the generated APIs to provide resumable uploads.
@@ -177,13 +182,18 @@ func (rx *ResumableUpload) Upload(ctx context.Context) (resp *http.Response, err
 		return resp, nil
 	}
 
+	chunkDeadline := defaultRetryDeadline
+	if retryDeadline, ok := ctx.Value(retryDeadlineKey).(time.Duration); ok {
+		chunkDeadline = retryDeadline
+	}
+
 	// Send all chunks.
 	for {
 		var pause time.Duration
 
 		// Each chunk gets its own initialized-at-zero retry.
 		bo := backoff()
-		quitAfter := time.After(retryDeadline)
+		quitAfter := time.After(chunkDeadline)
 
 		// Retry loop for a single chunk.
 		for {
diff --git a/internal/gensupport/resumable_test.go b/internal/gensupport/resumable_test.go
index 44544f7d71c..c2ec74b3e19 100644
--- a/internal/gensupport/resumable_test.go
+++ b/internal/gensupport/resumable_test.go
@@ -337,17 +337,13 @@ func TestRetry_EachChunkHasItsOwnRetryDeadline(t *testing.T) {
 		Callback:  func(int64) {},
 	}
 
-	oldRetryDeadline := retryDeadline
-	retryDeadline = 5 * time.Second
-	defer func() { retryDeadline = oldRetryDeadline }()
-
 	oldBackoff := backoff
 	backoff = func() Backoff { return new(PauseOneSecond) }
 	defer func() { backoff = oldBackoff }()
 
 	resCode := make(chan int, 1)
 	go func() {
-		resp, err := rx.Upload(context.Background())
+		resp, err := rx.Upload(context.WithValue(context.Background(), retryDeadlineKey, 5*time.Second))
 		if err != nil {
 			t.Error(err)
 			return