fix(spanner): Fix blind retry for ResourceExhausted (#12523)

sakthivelmanii · web-flow · commit f9b6e88bd3fc · 2025-07-08T11:38:57.000+05:30
* fix(spanner): Fix blind retry for ResourceExhausted

* Fix indentation issue in retry.go

* Fix comments
diff --git a/spanner/read_test.go b/spanner/read_test.go
@@ -30,9 +30,11 @@ import (
 	"github.com/googleapis/gax-go/v2"
 	"go.opentelemetry.io/otel/metric/noop"
 	"google.golang.org/api/iterator"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/durationpb"
 	proto3 "google.golang.org/protobuf/types/known/structpb"
 	structpb "google.golang.org/protobuf/types/known/structpb"
 )
@@ -1552,6 +1554,127 @@ func TestGrpcReconnect(t *testing.T) {
 	}
 }
 
+func TestRetryResourceExhaustedWithoutRetryInfo(t *testing.T) {
+	restore := setMaxBytesBetweenResumeTokens()
+	defer restore()
+
+	server, c, teardown := setupMockedTestServer(t)
+	defer teardown()
+	mc, err := c.sc.nextClient()
+	if err != nil {
+		t.Fatalf("failed to create a grpc client")
+	}
+
+	session, err := createSession(mc)
+	if err != nil {
+		t.Fatalf("failed to create a session")
+	}
+
+	// Simulate an ResourceExhausted error to interrupt the stream of PartialResultSet
+	// in order to test the grpc retrying mechanism.
+	server.TestSpanner.AddPartialResultSetError(
+		SelectSingerIDAlbumIDAlbumTitleFromAlbums,
+		PartialResultSetExecutionTime{
+			ResumeToken: EncodeResumeToken(2),
+			Err:         status.Errorf(codes.ResourceExhausted, "server is unavailable"),
+		},
+	)
+
+	// The retry is counted from the second call.
+	r := -1
+	// Establish a stream to mock cloud spanner server.
+	iter := stream(context.Background(), nil, c.metricsTracerFactory,
+		func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
+			r++
+			return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
+				Session:     session.Name,
+				Sql:         SelectSingerIDAlbumIDAlbumTitleFromAlbums,
+				ResumeToken: resumeToken,
+			}, opts...)
+
+		},
+		nil,
+		func(error) {}, mc.(*grpcSpannerClient))
+	defer iter.Stop()
+	for {
+		_, err := iter.Next()
+		if err == iterator.Done {
+			err = nil
+			break
+		}
+		if err != nil {
+			break
+		}
+	}
+	if r != 0 {
+		t.Errorf("retry count = %v, want 0", r)
+	}
+}
+
+// Verify that streaming query get retried upon ResourceExhausted real gRPC server
+// transport failures.
+func TestRetryResourceExhaustedWithRetryInfo(t *testing.T) {
+	restore := setMaxBytesBetweenResumeTokens()
+	defer restore()
+
+	server, c, teardown := setupMockedTestServer(t)
+	defer teardown()
+	mc, err := c.sc.nextClient()
+	if err != nil {
+		t.Fatalf("failed to create a grpc client")
+	}
+
+	session, err := createSession(mc)
+	if err != nil {
+		t.Fatalf("failed to create a session")
+	}
+
+	// Simulate an ResourceExhausted error to interrupt the stream of PartialResultSet
+	// in order to test the grpc retrying mechanism.
+	st := status.New(codes.ResourceExhausted, "server is unavailable")
+	retry := &errdetails.RetryInfo{
+		RetryDelay: durationpb.New(time.Nanosecond),
+	}
+	st, _ = st.WithDetails(retry)
+	server.TestSpanner.AddPartialResultSetError(
+		SelectSingerIDAlbumIDAlbumTitleFromAlbums,
+		PartialResultSetExecutionTime{
+			ResumeToken: EncodeResumeToken(2),
+			Err:         st.Err(),
+		},
+	)
+
+	// The retry is counted from the second call.
+	r := -1
+	// Establish a stream to mock cloud spanner server.
+	iter := stream(context.Background(), nil, c.metricsTracerFactory,
+		func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
+			r++
+			return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
+				Session:     session.Name,
+				Sql:         SelectSingerIDAlbumIDAlbumTitleFromAlbums,
+				ResumeToken: resumeToken,
+			}, opts...)
+
+		},
+		nil,
+		func(error) {}, mc.(*grpcSpannerClient))
+	defer iter.Stop()
+	for {
+		_, err := iter.Next()
+		if err == iterator.Done {
+			err = nil
+			break
+		}
+		if err != nil {
+			break
+		}
+	}
+	if r != 1 {
+		t.Errorf("retry count = %v, want 1", r)
+	}
+}
+
 // Test cancel/timeout for client operations.
 func TestCancelTimeout(t *testing.T) {
 	restore := setMaxBytesBetweenResumeTokens()
diff --git a/spanner/retry.go b/spanner/retry.go
@@ -59,7 +59,8 @@ func onCodes(bo gax.Backoff, cc ...codes.Code) gax.Retryer {
 // Retry returns the retry delay returned by Cloud Spanner if that is present.
 // Otherwise it returns the retry delay calculated by the generic gax Retryer.
 func (r *spannerRetryer) Retry(err error) (time.Duration, bool) {
-	if status.Code(err) == codes.Internal &&
+	errCode := status.Code(err)
+	if errCode == codes.Internal &&
 		!strings.Contains(err.Error(), "stream terminated by RST_STREAM") &&
 		// See b/25451313.
 		!strings.Contains(err.Error(), "HTTP/2 error code: INTERNAL_ERROR") &&
@@ -74,7 +75,14 @@ func (r *spannerRetryer) Retry(err error) (time.Duration, bool) {
 	if !shouldRetry {
 		return 0, false
 	}
-	if serverDelay, hasServerDelay := ExtractRetryDelay(err); hasServerDelay {
+
+	serverDelay, hasServerDelay := ExtractRetryDelay(err)
+	// Retry ResourceExhausted error only if there's a server delay in the trailer
+	if errCode == codes.ResourceExhausted && (!hasServerDelay || serverDelay <= 0) {
+		return 0, false
+	}
+
+	if hasServerDelay {
 		delay = serverDelay
 	}
 	return delay, true
