feat: adding read stall retry for gRPC storage client

storage/client_test.go

@@ -3295,7 +3295,7 @@ func TestRetryReadStallEmulated(t *testing.T) {
 	client, err := NewClient(ctx, experimental.WithReadStallTimeout(
 		&experimental.ReadStallTimeoutConfig{
 			TargetPercentile: 0.99,
-			Min:              10 * time.Millisecond,
+			Min:              250 * time.Millisecond,
 		}))
 	if err != nil {
 		t.Fatalf("storage.NewClient: %v", err)
@@ -3335,6 +3335,59 @@ func TestRetryReadStallEmulated(t *testing.T) {
 	}
 }
 
+func TestGRPCRetryReadStallEmulated(t *testing.T) {
+	checkEmulatorEnvironment(t)
+	t.Run("ReadObject", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		opts := []option.ClientOption{
+			experimental.WithReadStallTimeout(
+				&experimental.ReadStallTimeoutConfig{
+					TargetPercentile: 0.99,
+					Min:              250 * time.Millisecond,
+				}),
+		}
+
+		client, err := NewGRPCClient(ctx, opts...)
+		if err != nil {
+			t.Fatalf("storage.NewGRPCClient: %v", err)
+		}
+		defer client.Close()
+		client.SetRetry(WithBackoff(gax.Backoff{Initial: 10 * time.Millisecond}))
+
+		project := "fake-project"
+		bucket := fmt.Sprintf("grpc-bucket-%d", time.Now().Nanosecond())
+		if err := client.Bucket(bucket).Create(ctx, project, nil); err != nil {
+			t.Fatalf("client.Bucket.Create: %v", err)
+		}
+
+		name, _, _, err := createObjectWithContent(ctx, bucket, randomBytes3MiB)
+		if err != nil {
+			t.Fatalf("createObject: %v", err)
+		}
+
+		instructions := map[string][]string{"storage.objects.get": {"stall-for-10s-after-0K"}}
+		testID := createRetryTest(t, client.tc, instructions)
+
+		testCtx := callctx.SetHeaders(ctx, "x-retry-test-id", testID)
+
+		r, err := client.Bucket(bucket).Object(name).NewReader(testCtx)
+		if err != nil {
+			t.Fatalf("NewReader: %v", err)
+		}
+		defer r.Close()
+
+		buf := &bytes.Buffer{}
+		if _, err := io.Copy(buf, r); err != nil {
+			t.Fatalf("io.Copy: %v", err)
+		}
+		if !bytes.Equal(buf.Bytes(), randomBytes3MiB) {
+			t.Errorf("content does not match, got len %v, want len %v", buf.Len(), len(randomBytes3MiB))
+		}
+	})
+}
+
 func TestWriterChunkTransferTimeoutEmulated(t *testing.T) {
 	transportClientTest(skipGRPC("service is not implemented"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
 		_, err := client.CreateBucket(ctx, project, bucket, &BucketAttrs{}, nil)

storage/dynamic_delay.go

@@ -15,10 +15,14 @@
 package storage
 
 import (
+	"context"
 	"fmt"
+	"log"
 	"math"
 	"sync"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 // dynamicDelay dynamically calculates the delay at a fixed percentile, based on
@@ -235,3 +239,55 @@ func (b *bucketDelayManager) update(bucketName string, latency time.Duration) {
 func (b *bucketDelayManager) getValue(bucketName string) time.Duration {
 	return b.getDelay(bucketName).getValue()
 }
+
+// executeWithReadStallTimeout executes openStream with dynamic delay stall retry tracking.
+func executeWithReadStallTimeout(
+	ctx context.Context,
+	dm *bucketDelayManager,
+	bucket string,
+	requestID uuid.UUID,
+	openStream func(ctx context.Context) error,
+	onStall func(),
+) error {
+	if dm == nil {
+		return openStream(ctx)
+	}
+
+	cancelCtx, cancel := context.WithCancel(ctx)
+	var (
+		innerErr error
+		done     = make(chan struct{}, 1)
+	)
+
+	go func() {
+		reqStartTime := time.Now()
+		innerErr = openStream(cancelCtx)
+		if innerErr == nil {
+			reqLatency := time.Since(reqStartTime)
+			dm.update(bucket, reqLatency)
+		} else if ctx.Err() == nil && cancelCtx.Err() != nil {
+			dm.increase(bucket)
+		}
+		done <- struct{}{}
+	}()
+
+	stallTimeout := dm.getValue(bucket)
+	timer := time.NewTimer(stallTimeout)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		log.Printf("[%s] stalled read-req cancelled after %fs", requestID, stallTimeout.Seconds())
+		cancel()
+		<-done
+		if onStall != nil {
+			onStall()
+		}
+		return context.DeadlineExceeded
+	case <-done:
+		if innerErr != nil {
+			cancel()
+		}
+	}
+	return innerErr
+}

storage/dynamic_delay_test.go

@@ -13,6 +13,8 @@
 package storage
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"math"
 	"math/rand"
@@ -22,6 +24,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/google/uuid"
 )
 
 func applySamples(numSamples int, expectedValue float64, rnd *rand.Rand, d *dynamicDelay) int {
@@ -383,3 +386,75 @@ func TestBucketDelayManagerMapSize(t *testing.T) {
 		t.Errorf("Expected %d buckets in the map, but got %d", numBuckets, len(b.delays))
 	}
 }
+
+func TestExecuteWithReadStallTimeout(t *testing.T) {
+	requestID := uuid.New()
+
+	t.Run("nil_delay_manager", func(t *testing.T) {
+		executed := false
+		err := executeWithReadStallTimeout(context.Background(), nil, "bucket", requestID, func(ctx context.Context) error {
+			executed = true
+			return nil
+		}, nil)
+		if err != nil {
+			t.Errorf("expected nil error, got %v", err)
+		}
+		if !executed {
+			t.Error("expected openStream to be executed")
+		}
+	})
+
+	t.Run("fast_success", func(t *testing.T) {
+		dm, _ := newBucketDelayManager(0.99, 1.5, 100*time.Millisecond, 10*time.Millisecond, 10*time.Second)
+		err := executeWithReadStallTimeout(context.Background(), dm, "bucket", requestID, func(ctx context.Context) error {
+			time.Sleep(2 * time.Millisecond)
+			return nil
+		}, nil)
+		if err != nil {
+			t.Errorf("expected nil error, got %v", err)
+		}
+	})
+
+	t.Run("stall_timeout_triggered", func(t *testing.T) {
+		dm, _ := newBucketDelayManager(0.99, 1.5, 10*time.Millisecond, 10*time.Millisecond, 10*time.Second)
+		initialVal := dm.getValue("bucket")
+		stalled := false
+
+		err := executeWithReadStallTimeout(context.Background(), dm, "bucket", requestID, func(ctx context.Context) error {
+			<-ctx.Done()
+			return ctx.Err()
+		}, func() {
+			stalled = true
+		})
+
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Errorf("expected context.DeadlineExceeded, got %v", err)
+		}
+		if !stalled {
+			t.Error("expected onStall callback to be invoked")
+		}
+		if newVal := dm.getValue("bucket"); newVal <= initialVal {
+			t.Errorf("expected dynamic timeout to increase, initial %v, new %v", initialVal, newVal)
+		}
+	})
+
+	t.Run("outer_context_cancelled", func(t *testing.T) {
+		dm, _ := newBucketDelayManager(0.99, 1.5, 100*time.Millisecond, 10*time.Millisecond, 10*time.Second)
+		initialVal := dm.getValue("bucket")
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel() // Cancel outer context immediately.
+
+		err := executeWithReadStallTimeout(ctx, dm, "bucket", requestID, func(ctx context.Context) error {
+			<-ctx.Done()
+			return ctx.Err()
+		}, nil)
+
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("expected context.Canceled, got %v", err)
+		}
+		if newVal := dm.getValue("bucket"); newVal != initialVal {
+			t.Errorf("expected dynamic timeout NOT to increase on outer context cancellation, initial %v, new %v", initialVal, newVal)
+		}
+	})
+}

storage/experimental/experimental.go

@@ -52,17 +52,16 @@ func WithMeterProvider(mp *metric.MeterProvider) option.ClientOption {
 	return internal.WithMeterProvider.(func(*metric.MeterProvider) option.ClientOption)(mp)
 }
 
-// WithReadStallTimeout provides a [option.ClientOption] that may be passed to [storage.NewClient].
+// WithReadStallTimeout provides a [option.ClientOption] that may be passed to [storage.NewClient] or [storage.NewGRPCClient].
 // It enables the client to retry stalled requests when starting a download from
 // Cloud Storage. If the timeout elapses with no response from the server, the request
 // is automatically retried.
 // The timeout is initially set to ReadStallTimeoutConfig.Min. The client tracks
 // latency across all read requests from the client for each bucket accessed, and can
 // adjust the timeout higher to the target percentile when latency for request to that
 // bucket is high.
-// Currently, this is supported only for downloads ([storage.NewReader] and
-// [storage.NewRangeReader] calls) and only for the XML API. Other read APIs (gRPC & JSON)
-// will be supported soon.
+// Currently, this is supported for downloads ([storage.NewReader] and
+// [storage.NewRangeReader] calls) on the XML API and gRPC API (not implemented for JSON API).
 func WithReadStallTimeout(rstc *ReadStallTimeoutConfig) option.ClientOption {
 	return internal.WithReadStallTimeout.(func(config *ReadStallTimeoutConfig) option.ClientOption)(rstc)
 }

storage/grpc_client.go

@@ -29,6 +29,7 @@ import (
 	"cloud.google.com/go/iam/apiv1/iampb"
 	gapic "cloud.google.com/go/storage/internal/apiv2"
 	"cloud.google.com/go/storage/internal/apiv2/storagepb"
+	"github.com/google/uuid"
 	"github.com/googleapis/gax-go/v2"
 	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
@@ -120,12 +121,13 @@ func defaultGRPCOptions() []option.ClientOption {
 // grpcStorageClient is the gRPC API implementation of the transport-agnostic
 // storageClient interface.
 type grpcStorageClient struct {
-	raw            *gapic.Client
-	settings       *settings
-	config         *storageConfig
-	dpDiag         string
-	metrics        *clientMetrics
-	metricsCleanup func()
+	raw                        *gapic.Client
+	settings                   *settings
+	config                     *storageConfig
+	dynamicReadReqStallTimeout *bucketDelayManager
+	dpDiag                     string
+	metrics                    *clientMetrics
+	metricsCleanup             func()
 
 	// configFeatureAttributes tracks client-level features that are enabled for this
 	// client instance.
@@ -200,11 +202,26 @@ func newGRPCStorageClient(ctx context.Context, opts ...storageOption) (client *g
 		}()
 	}
 
+	var bd *bucketDelayManager
+	if config.readStallTimeoutConfig != nil {
+		drrstConfig := config.readStallTimeoutConfig
+		bd, err = newBucketDelayManager(
+			drrstConfig.TargetPercentile,
+			getDynamicReadReqIncreaseRateFromEnv(),
+			getDynamicReadReqInitialTimeoutSecFromEnv(drrstConfig.Min),
+			drrstConfig.Min,
+			defaultDynamicReqdReqMaxTimeout)
+		if err != nil {
+			return nil, fmt.Errorf("creating dynamic-delay: %w", err)
+		}
+	}
+
 	c := &grpcStorageClient{
-		settings:       s,
-		config:         &config,
-		metrics:        clientMetrics,
-		metricsCleanup: metricsCleanup,
+		settings:                   s,
+		config:                     &config,
+		metrics:                    clientMetrics,
+		metricsCleanup:             metricsCleanup,
+		dynamicReadReqStallTimeout: bd,
 	}
 	// Add routing interceptors to inject headers.
 	ui, si := c.routingInterceptors()
@@ -1251,6 +1268,7 @@ func (c *grpcStorageClient) NewRangeReader(ctx context.Context, params *newRange
 	ctx, _ = startSpan(ctx, "grpcStorageClient.NewRangeReader")
 	defer func() { endSpan(ctx, err) }()
 
+	requestID := uuid.New()
 	s := callSettings(c.settings, opts...)
 
 	s.gax = append(s.gax, gax.WithGRPCOptions(
@@ -1310,7 +1328,7 @@ func (c *grpcStorageClient) NewRangeReader(ctx context.Context, params *newRange
 		var err error
 		var decoder *readResponseDecoder
 
-		err = run(cc, func(ctx context.Context) error {
+		openStream := func(ctx context.Context) error {
 			var databufs mem.BufferSlice
 			openAndSendReq := func() error {
 				databufs = mem.BufferSlice{}
@@ -1368,6 +1386,15 @@ func (c *grpcStorageClient) NewRangeReader(ctx context.Context, params *newRange
 			}
 			err = decoder.readFullObjectResponse()
 			return err
+		}
+
+		err = run(cc, func(ctx context.Context) error {
+			return executeWithReadStallTimeout(ctx, c.dynamicReadReqStallTimeout, params.bucket, requestID, openStream, func() {
+				if decoder != nil && decoder.databufs != nil {
+					decoder.databufs.Free()
+					decoder = nil
+				}
+			})
 		}, s.retry, s.idempotent, withOperation("ReadObject"), withBucket(params.bucket), withObject(params.object))
 		if err != nil {
 			// Close the stream context we just created to ensure we don't leak

storage/grpc_reader.go

@@ -23,6 +23,7 @@ import (
 	"io"
 
 	"cloud.google.com/go/storage/internal/apiv2/storagepb"
+	"github.com/google/uuid"
 	"github.com/googleapis/gax-go/v2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/encoding"
@@ -85,6 +86,7 @@ func (c *grpcStorageClient) NewRangeReaderReadObject(ctx context.Context, params
 	ctx, _ = startSpan(ctx, "grpcStorageClient.NewRangeReaderReadObject")
 	defer func() { endSpan(ctx, err) }()
 
+	requestID := uuid.New()
 	s := callSettings(c.settings, opts...)
 
 	s.gax = append(s.gax, gax.WithGRPCOptions(
@@ -134,7 +136,7 @@ func (c *grpcStorageClient) NewRangeReaderReadObject(ctx context.Context, params
 		var err error
 		var decoder *readObjectResponseDecoder
 
-		err = run(cc, func(ctx context.Context) error {
+		openStream := func(ctx context.Context) error {
 			stream, err = c.raw.ReadObject(ctx, req, s.gax...)
 			if err != nil {
 				return err
@@ -157,6 +159,15 @@ func (c *grpcStorageClient) NewRangeReaderReadObject(ctx context.Context, params
 			}
 			err = decoder.readFullObjectResponse()
 			return err
+		}
+
+		err = run(cc, func(ctx context.Context) error {
+			return executeWithReadStallTimeout(ctx, c.dynamicReadReqStallTimeout, params.bucket, requestID, openStream, func() {
+				if decoder != nil && decoder.databufs != nil {
+					decoder.databufs.Free()
+					decoder = nil
+				}
+			})
 		}, s.retry, s.idempotent, withOperation("ReadObject"), withBucket(params.bucket), withObject(params.object))
 		if err != nil {
 			// Close the stream context we just created to ensure we don't leak