feat(storage): Add stall timeouts to first chunk in JSON writes [WIP]

google-api-go-generator/gen.go

@@ -826,6 +826,7 @@ func (a *API) GenerateCode() ([]byte, error) {
 	pn("")
 	if a.Name == "storage" {
 		pn("  %q", "github.com/googleapis/gax-go/v2")
+		pn("  %q", "time")
 	}
 	pn("  %q", "github.com/googleapis/gax-go/v2/internallog")
 	for _, imp := range []struct {
@@ -860,6 +861,7 @@ func (a *API) GenerateCode() ([]byte, error) {
 	pn("var _ = internallog.New")
 	if a.Name == "storage" {
 		pn("var _ = gax.Version")
+		pn("var _ = time.Second")
 	}
 	pn("")
 	pn("const apiId = %q", a.doc.ID)
@@ -2222,6 +2224,16 @@ func (meth *Method) generateCode() {
 		pn("	}")
 		pn("	return c")
 		pn("}")
+		commentStall := "WithStallTimeout sets a per-attempt/request timeout. If a request " +
+			"attempt takes longer than this duration, it is cancelled and retried if retryable."
+		p("\n%s", asComment("", commentStall))
+		pn("func (c *%s) WithStallTimeout(timeout time.Duration) *%s {", callName, callName)
+		pn("	if c.retry == nil {")
+		pn("		c.retry = &gensupport.RetryConfig{}")
+		pn("	}")
+		pn("	c.retry.RequestTimeout = timeout")
+		pn("	return c")
+		pn("}")
 	}
 
 	comment := "Fields allows partial responses to be retrieved. See " +

internal/gensupport/retry.go

@@ -90,6 +90,9 @@ func shouldRetry(status int, err error) bool {
 type RetryConfig struct {
 	Backoff     *gax.Backoff
 	ShouldRetry func(err error) bool
+	// RequestTimeout is a per-attempt/request timeout. If a request attempt takes
+	// longer than this duration, it is cancelled and retried if retryable.
+	RequestTimeout time.Duration
 }
 
 // Get a new backoff object based on the configured values.

internal/gensupport/send.go

@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"strings"
 	"time"
@@ -132,6 +133,19 @@ func SendRequestWithRetry(ctx context.Context, client *http.Client, req *http.Re
 	return sendAndRetry(ctx, client, req, retry)
 }
 
+type cancelCloser struct {
+	io.ReadCloser
+	cancel context.CancelFunc
+}
+
+func (c *cancelCloser) Close() error {
+	err := c.ReadCloser.Close()
+	if c.cancel != nil {
+		c.cancel()
+	}
+	return err
+}
+
 func sendAndRetry(ctx context.Context, client *http.Client, req *http.Request, retry *RetryConfig) (*http.Response, error) {
 	if client == nil {
 		client = http.DefaultClient
@@ -193,17 +207,42 @@ func sendAndRetry(ctx context.Context, client *http.Client, req *http.Request, r
 		req.Header.Set("X-Goog-Api-Client", xGoogHeader)
 		req.Header.Set("X-Goog-Gcs-Idempotency-Token", invocationID)
 
-		resp, err = client.Do(req.WithContext(ctx))
+		var reqCtx context.Context
+		var cancel context.CancelFunc
+		stallRetryAvailable := false
+		if retry != nil &&
+			retry.RequestTimeout > 0 &&
+			req.GetBody != nil {
+			stallRetryAvailable = true
+			reqCtx, cancel = context.WithTimeout(ctx, retry.RequestTimeout)
+		} else {
+			reqCtx = ctx
+		}
+
+		resp, err = client.Do(req.WithContext(reqCtx))
 
 		var status int
 		if resp != nil {
 			status = resp.StatusCode
 		}
 
+		isStallTimeout := false
+		if stallRetryAvailable &&
+			ctx.Err() == nil &&
+			(errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled)) {
+			log.Printf("Cancelled the request due to stall")
+			isStallTimeout = true
+		}
+
+		if resp != nil && resp.Body != nil && cancel != nil {
+			resp.Body = &cancelCloser{ReadCloser: resp.Body, cancel: cancel}
+		} else if cancel != nil {
+			cancel()
+		}
 		// Check if we can retry the request. A retry can only be done if the error
 		// is retryable and the request body can be re-created using GetBody (this
 		// will not be possible if the body was unbuffered).
-		if req.GetBody == nil || !errorFunc(status, err) {
+		if req.GetBody == nil || (!isStallTimeout && !errorFunc(status, err)) {
 			break
 		}
 		attempts++

internal/gensupport/send_test.go

@@ -8,11 +8,15 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"regexp"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/googleapis/gax-go/v2"
 	"github.com/googleapis/gax-go/v2/callctx"
 )
 
@@ -138,3 +142,49 @@ func TestCanceledContextDoesNotPerformRequest(t *testing.T) {
 		}
 	}
 }
+
+type stubRoundTripper struct {
+	roundTrip func(*http.Request) (*http.Response, error)
+}
+
+func (rt *stubRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	return rt.roundTrip(r)
+}
+
+func TestSendRequestWithRetry_RequestTimeout(t *testing.T) {
+	attempts := 0
+	transport := &stubRoundTripper{
+		roundTrip: func(r *http.Request) (*http.Response, error) {
+			attempts++
+			if attempts < 3 {
+				// Simulate stall by blocking until the request context is cancelled.
+				<-r.Context().Done()
+				return nil, r.Context().Err()
+			}
+			return &http.Response{StatusCode: 200}, nil
+		},
+	}
+	client := http.Client{Transport: transport}
+	req, _ := http.NewRequest("GET", "url", nil)
+	req.GetBody = func() (io.ReadCloser, error) {
+		return io.NopCloser(strings.NewReader("")), nil
+	}
+
+	retry := &RetryConfig{
+		RequestTimeout: 50 * time.Millisecond,
+		Backoff: &gax.Backoff{
+			Initial: 1 * time.Millisecond,
+		},
+	}
+
+	resp, err := SendRequestWithRetry(context.Background(), &client, req, retry)
+	if err != nil {
+		t.Fatalf("SendRequestWithRetry failed: %v", err)
+	}
+	if resp.StatusCode != 200 {
+		t.Errorf("got status %d, want 200", resp.StatusCode)
+	}
+	if attempts != 3 {
+		t.Errorf("got %d attempts, want 3", attempts)
+	}
+}