From 498a746769fa43958b92af8875b927879947128e Mon Sep 17 00:00:00 2001
From: Chris Cotter <cjcotter@google.com>
Date: Thu, 10 Feb 2022 18:09:50 -0500
Subject: [PATCH] feat(storage): add Writer.ChunkRetryDeadline (#5482)

Expose the ChunkRetryDeadline MediaOption through the manual
client layer. This allows users to set a longer deadline for
chunk retries in resumable uploads if desired.

I also considered this as a RetryOption, but 1. it's only
relevant for uploads and 2. it should be configured in
conjunction with ChunkSize which is a field on Writer.

Updates googleapis/google-api-go-client#685
---
 storage/retry_test.go | 13 ++++++-------
 storage/writer.go     | 21 +++++++++++++++++++++
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/storage/retry_test.go b/storage/retry_test.go
index fe66e443f7c..7203e5ce7dc 100644
--- a/storage/retry_test.go
+++ b/storage/retry_test.go
@@ -33,9 +33,6 @@ import (
 )
 
 func TestIndefiniteRetries(t *testing.T) {
-	if testing.Short() {
-		t.Skip("A long running test for retries")
-	}
 
 	uploadRoute := "/upload"
 
@@ -138,6 +135,9 @@ func TestIndefiniteRetries(t *testing.T) {
 	maxFileSize := 1 << 20
 	w.ChunkSize = maxFileSize / 4
 
+	// Use a shorter retry deadline to speed up the test.
+	w.ChunkRetryDeadline = time.Second
+
 	for i := 0; i < maxFileSize; {
 		nowStr := time.Now().Format(time.RFC3339Nano)
 		n, err := fmt.Fprintf(w, "%s", nowStr)
@@ -153,10 +153,9 @@ func TestIndefiniteRetries(t *testing.T) {
 		closeDone <- w.Close()
 	}()
 
-	// Given that the ExponentialBackoff is 30 seconds from a start of 100ms,
-	// let's wait for a maximum of 5 minutes to account for (2**n) increments
-	// between [100ms, 30s].
-	maxWait := 5 * time.Minute
+	// Given the exponential backoff math and the ChunkRetryDeadline math, use a
+	// max value of 10s. Experimentally this test usually passes in < 5s.
+	maxWait := 10 * time.Second
 	select {
 	case <-time.After(maxWait):
 		t.Fatalf("Test took longer than %s to return", maxWait)
diff --git a/storage/writer.go b/storage/writer.go
index 51d8b63779b..bd886b04e01 100644
--- a/storage/writer.go
+++ b/storage/writer.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io"
 	"sync"
+	"time"
 	"unicode/utf8"
 
 	"github.com/golang/protobuf/proto"
@@ -76,6 +77,23 @@ type Writer struct {
 	// ChunkSize must be set before the first Write call.
 	ChunkSize int
 
+	// ChunkRetryDeadline sets a per-chunk retry deadline for multi-chunk
+	// resumable uploads.
+	//
+	// For uploads of larger files, the Writer will attempt to retry if the
+	// request to upload a particular chunk fails with a transient error.
+	// If a single chunk has been attempting to upload for longer than this
+	// deadline and the request fails, it will no longer be retried, and the error
+	// will be returned to the caller. This is only applicable for files which are
+	// large enough to require a multi-chunk resumable upload. The default value
+	// is 32s. Users may want to pick a longer deadline if they are using larger
+	// values for ChunkSize or if they expect to have a slow or unreliable
+	// internet connection.
+	//
+	// To set a deadline on the entire upload, use context timeout or
+	// cancellation.
+	ChunkRetryDeadline time.Duration
+
 	// ProgressFunc can be used to monitor the progress of a large write.
 	// operation. If ProgressFunc is not nil and writing requires multiple
 	// calls to the underlying service (see
@@ -127,6 +145,9 @@ func (w *Writer) open() error {
 	if c := attrs.ContentType; c != "" {
 		mediaOpts = append(mediaOpts, googleapi.ContentType(c))
 	}
+	if w.ChunkRetryDeadline != 0 {
+		mediaOpts = append(mediaOpts, googleapi.ChunkRetryDeadline(w.ChunkRetryDeadline))
+	}
 
 	go func() {
 		defer close(w.donec)