Changes to gRPC message too large error handling (#2042)

* Changed gRPC message too large error handling to wrap error in a new type instead of using non-standard cause enum

* Improved TestGrpcMessageTooLarge to not rely on timeouts

Author: maciejdudko

internal/common/retry/interceptor.go

@@ -9,8 +9,6 @@ import (
 
 	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/util/backoffutils"
-	enumspb "go.temporal.io/api/enums/v1"
-	"go.temporal.io/api/errordetails/v1"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -45,6 +43,11 @@ type (
 		maximumAttempts    int
 	}
 
+	GrpcMessageTooLargeError struct {
+		err    error
+		status *status.Status
+	}
+
 	contextKey struct{}
 )
 
@@ -94,15 +97,15 @@ func NewGrpcRetryConfig(initialInterval time.Duration) *GrpcRetryConfig {
 var (
 	// ConfigKey context key for GrpcRetryConfig
 	ConfigKey = contextKey{}
-	// gRPC response codes that represent unconditionally retryable errors.
+	// gRPC response codes that represent retryable errors.
 	// The following status codes are never retried by the library:
 	//    INVALID_ARGUMENT, NOT_FOUND, ALREADY_EXISTS, FAILED_PRECONDITION, ABORTED, OUT_OF_RANGE, DATA_LOSS
 	// codes.DeadlineExceeded and codes.Canceled are not here (and shouldn't be here!)
 	// because they are coming from go context and "context errors are not retriable based on user settings"
 	// by gRPC library.
-	// codes.Internal and codes.ResourceExhausted have special logic for whether they are retryable or not,
-	// and so they're not included in this list.
-	alwaysRetryableCodes = []codes.Code{codes.Aborted, codes.Unavailable, codes.Unknown}
+	// codes.ResourceExhausted is non-retryable if it comes from GrpcMessageTooLargeError, but otherwise is retryable.
+	// codes.Internal is not included because it's retryable or non-retryable depending on server capabilities.
+	retryableCodesWithoutInternal = []codes.Code{codes.Aborted, codes.ResourceExhausted, codes.Unavailable, codes.Unknown}
 )
 
 // NewRetryOptionsInterceptor creates a new gRPC interceptor that populates retry options for each call based on values
@@ -132,7 +135,7 @@ func NewRetryOptionsInterceptor(excludeInternal *atomic.Bool) grpc.UnaryClientIn
 				opts = append(opts, grpc_retry.WithMax(math.MaxUint32))
 			}
 			opts = append(opts, grpc_retry.WithRetriable(func(err error) bool {
-				return IsRetryable(status.Convert(err), excludeInternal)
+				return IsRetryable(err, excludeInternal)
 			}))
 		} else {
 			// Do not retry if retry config is not set.
@@ -142,51 +145,48 @@ func NewRetryOptionsInterceptor(excludeInternal *atomic.Bool) grpc.UnaryClientIn
 	}
 }
 
-func IsRetryable(status *status.Status, excludeInternalFromRetry *atomic.Bool) bool {
-	if status == nil {
+func IsRetryable(err error, excludeInternalFromRetry *atomic.Bool) bool {
+	if _, ok := err.(*GrpcMessageTooLargeError); ok {
+		return false
+	}
+	grpcStatus := status.Convert(err)
+	if grpcStatus == nil {
 		return false
 	}
-	errCode := status.Code()
-	for _, retryable := range alwaysRetryableCodes {
+	errCode := grpcStatus.Code()
+	for _, retryable := range retryableCodesWithoutInternal {
 		if errCode == retryable {
 			return true
 		}
 	}
 	if errCode == codes.Internal {
 		return !excludeInternalFromRetry.Load()
 	}
-	if errCode == codes.ResourceExhausted {
-		if details := status.Details(); len(details) > 0 {
-			if failure, ok := details[0].(*errordetails.ResourceExhaustedFailure); ok {
-				return failure.Cause != RESOURCE_EXHAUSTED_CAUSE_EXT_GRPC_MESSAGE_TOO_LARGE
-			}
-		}
-	}
 	return false
 }
 
-// RESOURCE_EXHAUSTED_CAUSE_EXT_GRPC_MESSAGE_TOO_LARGE is an extension to the ResourceExhaustedCause enum to mark gRPC message too large errors.
-const RESOURCE_EXHAUSTED_CAUSE_EXT_GRPC_MESSAGE_TOO_LARGE enumspb.ResourceExhaustedCause = 101 // TODO: add the cause to the upstream API repo and remove this (see https://github.com/temporalio/sdk-go/issues/2030)
-
-// SetGrpcMessageTooLargeErrorCauseInterceptor adds appropriate error details if the error cause is gRPC message being too large.
-func SetGrpcMessageTooLargeErrorCauseInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+// GrpcMessageTooLargeErrorInterceptor checks if the error is caused by gRPC message being too large and converts it into GrpcMessageTooLargeError.
+func GrpcMessageTooLargeErrorInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 	err := invoker(ctx, method, req, reply, cc, opts...)
-	if grpcStatus := status.Convert(err); isGrpcMessageTooLargeError(grpcStatus) {
-		// Copying code and message but ignoring original details
-		newStatus := status.New(grpcStatus.Code(), grpcStatus.Message())
-		newStatus, detailsErr := newStatus.WithDetails(&errordetails.ResourceExhaustedFailure{
-			Cause: RESOURCE_EXHAUSTED_CAUSE_EXT_GRPC_MESSAGE_TOO_LARGE,
-		})
-		if detailsErr == nil {
-			if newErr := newStatus.Err(); newErr != nil {
-				err = newErr
-			}
-		}
+	if grpcStatus := status.Convert(err); isGrpcMessageTooLargeStatus(grpcStatus) {
+		err = &GrpcMessageTooLargeError{err: err, status: grpcStatus}
 	}
 	return err
 }
 
-func isGrpcMessageTooLargeError(status *status.Status) bool {
+func (e *GrpcMessageTooLargeError) Error() string {
+	return e.err.Error()
+}
+
+func (e *GrpcMessageTooLargeError) Unwrap() error {
+	return e.err
+}
+
+func (e *GrpcMessageTooLargeError) GRPCStatus() *status.Status {
+	return e.status
+}
+
+func isGrpcMessageTooLargeStatus(status *status.Status) bool {
 	if status == nil {
 		return false
 	}

internal/grpc_dialer.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"errors"
 	"sync/atomic"
 	"time"
 
@@ -139,7 +140,7 @@ func requiredInterceptors(
 		// Performs retries *IF* retry options are set for the call.
 		grpc_retry.UnaryClientInterceptor(),
 		// Prevents retrying grpc message too large errors, while allowing retries of other resource exhausted errors.
-		retry.SetGrpcMessageTooLargeErrorCauseInterceptor,
+		retry.GrpcMessageTooLargeErrorInterceptor,
 		// Report metrics for every call made to the server.
 		metrics.NewGRPCInterceptor(clientOptions.MetricsHandler, attemptSuffix, clientOptions.DisableErrorCodeMetricTags),
 	}
@@ -199,6 +200,9 @@ func headersProviderInterceptor(headersProvider HeadersProvider) grpc.UnaryClien
 
 func errorInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 	err := invoker(ctx, method, req, reply, cc, opts...)
-	err = serviceerror.FromStatus(status.Convert(err))
+	var grpcMessageTooLargeErr *retry.GrpcMessageTooLargeError
+	if !errors.As(err, &grpcMessageTooLargeErr) {
+		err = serviceerror.FromStatus(status.Convert(err))
+	}
 	return err
 }

internal/internal_task_handlers.go

@@ -24,7 +24,6 @@ import (
 	"go.temporal.io/api/serviceerror"
 	taskqueuepb "go.temporal.io/api/taskqueue/v1"
 	"go.temporal.io/api/workflowservice/v1"
-	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/durationpb"
 
 	"go.temporal.io/sdk/internal/common/retry"
@@ -2164,8 +2163,7 @@ func (i *temporalInvoker) internalHeartBeat(ctx context.Context, details *common
 		// Transient errors are getting retried for the duration of the heartbeat timeout.
 		// The fact that error has been returned means that activity should now be timed out, hence we should
 		// propagate cancellation to the handler.
-		statusErr, _ := status.FromError(err)
-		if retry.IsRetryable(statusErr, i.excludeInternalFromRetry) {
+		if retry.IsRetryable(err, i.excludeInternalFromRetry) {
 			i.cancelHandler(err)
 			isActivityCanceled = true
 		}

internal/internal_task_pollers.go

@@ -544,7 +544,8 @@ func (wtp *workflowTaskProcessor) RespondTaskCompletedWithMetrics(
 
 	response, err = wtp.sendTaskCompletedRequest(completedRequest, task)
 
-	if isGrpcMessageTooLargeError(err) {
+	var grpcMessageTooLargeErr *retry.GrpcMessageTooLargeError
+	if errors.As(err, &grpcMessageTooLargeErr) {
 		secondEmitFailMetric, secondErr := wtp.reportGrpcMessageTooLarge(completedRequest, task, err)
 		if secondEmitFailMetric {
 			emitFailMetric = true
@@ -1537,8 +1538,3 @@ func (nt *nexusTask) scaleDecision() (pollerScaleDecision, bool) {
 		pollRequestDeltaSuggestion: int(nt.task.PollerScalingDecision.PollRequestDeltaSuggestion),
 	}, true
 }
-
-func isGrpcMessageTooLargeError(err error) bool {
-	serviceErr, ok := err.(*serviceerror.ResourceExhausted)
-	return ok && serviceErr.Cause == retry.RESOURCE_EXHAUSTED_CAUSE_EXT_GRPC_MESSAGE_TOO_LARGE
-}

test/integration_test.go

@@ -7751,23 +7751,17 @@ func (ts *IntegrationTestSuite) TestLocalActivitySummary() {
 }
 
 func (ts *IntegrationTestSuite) TestGrpcMessageTooLarge() {
-	ts.T().Skip("issue-2033: Test is flaky")
-	assertGrpcErrorInHistoryOnce := func(ctx context.Context, run client.WorkflowRun) {
-		found := false
+	assertGrpcErrorInHistory := func(ctx context.Context, run client.WorkflowRun) {
 		iter := ts.client.GetWorkflowHistory(ctx, run.GetID(), run.GetRunID(), true, enumspb.HISTORY_EVENT_FILTER_TYPE_ALL_EVENT)
 		for iter.HasNext() {
 			event, err := iter.Next()
 			ts.NoError(err)
 			if event.EventType == enumspb.EVENT_TYPE_WORKFLOW_TASK_FAILED {
-				if found {
-					ts.Fail("Found more than 1 workflow task failed event in history")
-				} else {
-					found = true
-				}
 				ts.Equal(enumspb.WORKFLOW_TASK_FAILED_CAUSE_GRPC_MESSAGE_TOO_LARGE, event.GetWorkflowTaskFailedEventAttributes().Cause)
+				return
 			}
 		}
-		ts.True(found, "Workflow task failed event not found in history")
+		ts.Fail("Workflow task failed event not found in history")
 	}
 
 	veryLargeData := strings.Repeat("Very Large Data ", 500_000) // circa 8MB, double the default 4MB limit
@@ -7801,25 +7795,22 @@ func (ts *IntegrationTestSuite) TestGrpcMessageTooLarge() {
 	ts.worker.RegisterWorkflow(failureInQueryTaskWorkflowFn)
 	ts.worker.RegisterActivity(activityFn)
 	startOptions := client.StartWorkflowOptions{
-		TaskQueue:           ts.taskQueueName,
-		WorkflowTaskTimeout: 300 * time.Millisecond,
-		WorkflowRunTimeout:  1 * time.Second,
+		TaskQueue: ts.taskQueueName,
 	}
 
-	ts.Run("Activity start too large", func() {
-		run, err := ts.client.ExecuteWorkflow(ctx, startOptions, failureInWorkflowTaskWorkflowFn, true)
+	testFailureInWorkflowTask := func(success bool) {
+		run, err := ts.client.ExecuteWorkflow(ctx, startOptions, failureInWorkflowTaskWorkflowFn, success)
 		ts.NoError(err)
-		err = run.Get(ctx, nil)
-		ts.Error(err)
-		assertGrpcErrorInHistoryOnce(ctx, run)
+		assertGrpcErrorInHistory(ctx, run)
+		ts.NoError(ts.client.TerminateWorkflow(ctx, run.GetID(), run.GetRunID(), "Test completed"))
+	}
+
+	ts.Run("Activity start too large", func() {
+		testFailureInWorkflowTask(true)
 	})
 
 	ts.Run("Workflow failure too large", func() {
-		run, err := ts.client.ExecuteWorkflow(ctx, startOptions, failureInWorkflowTaskWorkflowFn, true)
-		ts.NoError(err)
-		err = run.Get(ctx, nil)
-		ts.Error(err)
-		assertGrpcErrorInHistoryOnce(ctx, run)
+		testFailureInWorkflowTask(false)
 	})
 
 	// successful query case is tested by TestLargeQueryResultError