Cancel History long polls before deadline

Author: stephanos

common/contextutil/timeout.go

@@ -0,0 +1,44 @@
+package contextutil
+
+import (
+	"context"
+	"time"
+)
+
+// WithDeadlineBuffer returns a child context whose deadline is capped by both:
+//   - the parent's deadline minus `buffer`
+//   - or, the given `maxTimeout` duration from now
+func WithDeadlineBuffer(
+	parent context.Context,
+	maxTimeout time.Duration,
+	buffer time.Duration,
+) (context.Context, context.CancelFunc) {
+	noop := func() {}
+
+	// If the parent is already done, don't bother creating a child.
+	if parent.Err() != nil {
+		return parent, noop
+	}
+
+	deadline, hasDeadline := parent.Deadline()
+	if !hasDeadline {
+		if maxTimeout == 0 {
+			return parent, noop
+		}
+		return context.WithTimeout(parent, maxTimeout)
+	}
+
+	// Inherit the parent's deadline if caller doesn't set a cap.
+	if maxTimeout == 0 {
+		return parent, noop
+	}
+
+	remaining := time.Until(deadline) - buffer
+	if remaining <= 0 {
+		return parent, noop
+	}
+	if remaining < maxTimeout {
+		maxTimeout = remaining
+	}
+	return context.WithTimeout(parent, maxTimeout)
+}

common/metrics/grpc_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/suite"
 	metricsspb "go.temporal.io/server/api/metrics/v1"
 	"go.temporal.io/server/common/log"
+	"go.temporal.io/server/common/testing/rpctest"
 	"go.uber.org/mock/gomock"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
@@ -36,7 +37,7 @@ func (s *grpcSuite) TearDownTest() {}
 func (s *grpcSuite) TestMetadataMetricInjection() {
 	logger := log.NewMockLogger(s.controller)
 	ctx := context.Background()
-	ssts := newMockServerTransportStream()
+	ssts := rpctest.NewMockServerTransportStream("/temporal.test/MetadataMetricInjection")
 	ctx = grpc.NewContextWithServerTransportStream(ctx, ssts)
 	anyMetricName := "any_metric_name"
 
@@ -73,8 +74,9 @@ func (s *grpcSuite) TestMetadataMetricInjection() {
 			)
 
 			s.Nil(err)
-			s.Equal(len(ssts.trailers), 1)
-			propagationContextBlobs := ssts.trailers[0].Get(metricsTrailerKey)
+			trailers := ssts.CapturedTrailers()
+			s.Equal(1, len(trailers))
+			propagationContextBlobs := trailers[0].Get(metricsTrailerKey)
 			s.NotNil(propagationContextBlobs)
 			s.Equal(1, len(propagationContextBlobs))
 			baggage := &metricsspb.Baggage{}
@@ -93,7 +95,7 @@ func (s *grpcSuite) TestMetadataMetricInjection() {
 func (s *grpcSuite) TestMetadataMetricInjection_NoMetricPresent() {
 	logger := log.NewMockLogger(s.controller)
 	ctx := context.Background()
-	ssts := newMockServerTransportStream()
+	ssts := rpctest.NewMockServerTransportStream("/temporal.test/MetadataMetricInjectionNoMetric")
 	ctx = grpc.NewContextWithServerTransportStream(ctx, ssts)
 
 	smcii := NewServerMetricsContextInjectorInterceptor()
@@ -128,8 +130,9 @@ func (s *grpcSuite) TestMetadataMetricInjection_NoMetricPresent() {
 			)
 
 			s.Nil(err)
-			s.Equal(len(ssts.trailers), 1)
-			propagationContextBlobs := ssts.trailers[0].Get(metricsTrailerKey)
+			trailers := ssts.CapturedTrailers()
+			s.Equal(1, len(trailers))
+			propagationContextBlobs := trailers[0].Get(metricsTrailerKey)
 			s.NotNil(propagationContextBlobs)
 			s.Equal(1, len(propagationContextBlobs))
 			baggage := &metricsspb.Baggage{}
@@ -162,26 +165,3 @@ func (s *grpcSuite) TestContextCounterAddNoMetricsContext() {
 	testCounterName := "test_counter"
 	ContextCounterAdd(context.Background(), testCounterName, 3)
 }
-
-func newMockServerTransportStream() *mockServerTransportStream {
-	return &mockServerTransportStream{trailers: []*metadata.MD{}}
-}
-
-type mockServerTransportStream struct {
-	trailers []*metadata.MD
-}
-
-func (s *mockServerTransportStream) Method() string {
-	return "mockssts"
-}
-func (s *mockServerTransportStream) SetHeader(md metadata.MD) error {
-	return nil
-}
-func (s *mockServerTransportStream) SendHeader(md metadata.MD) error {
-	return nil
-}
-func (s *mockServerTransportStream) SetTrailer(md metadata.MD) error {
-	mdCopy := md.Copy()
-	s.trailers = append(s.trailers, &mdCopy)
-	return nil
-}

common/rpc/grpc.go

@@ -3,21 +3,17 @@ package rpc
 import (
 	"context"
 	"crypto/tls"
-	"errors"
 	"time"
 
-	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
-	"go.temporal.io/server/common/log/tag"
 	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/common/rpc/interceptor"
 	serviceerrors "go.temporal.io/server/common/serviceerror"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/backoff"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
@@ -44,14 +40,6 @@ const (
 
 	// maxInternodeRecvPayloadSize indicates the internode max receive payload size.
 	maxInternodeRecvPayloadSize = 128 * 1024 * 1024 // 128 Mb
-
-	// ResourceExhaustedCauseHeader will be added to rpc response if request returns ResourceExhausted error.
-	// Value of this header will be ResourceExhaustedCause.
-	ResourceExhaustedCauseHeader = "X-Resource-Exhausted-Cause"
-
-	// ResourceExhaustedScopeHeader will be added to rpc response if request returns ResourceExhausted error.
-	// Value of this header will be the scope of exhausted resource.
-	ResourceExhaustedScopeHeader = "X-Resource-Exhausted-Scope"
 )
 
 // Dial creates a client connection to the given target with default options.
@@ -121,46 +109,3 @@ func headersInterceptor(
 	ctx = headers.Propagate(ctx)
 	return invoker(ctx, method, req, reply, cc, opts...)
 }
-
-func NewFrontendServiceErrorInterceptor(
-	logger log.Logger,
-) grpc.UnaryServerInterceptor {
-	return func(
-		ctx context.Context,
-		req interface{},
-		_ *grpc.UnaryServerInfo,
-		handler grpc.UnaryHandler,
-	) (interface{}, error) {
-
-		resp, err := handler(ctx, req)
-
-		if err == nil {
-			return resp, err
-		}
-
-		// mask some internal service errors at frontend
-		switch err.(type) {
-		case *serviceerrors.ShardOwnershipLost:
-			err = serviceerror.NewUnavailable("shard unavailable, please backoff and retry")
-		case *serviceerror.DataLoss:
-			err = serviceerror.NewUnavailable("internal history service error")
-		}
-
-		addHeadersForResourceExhausted(ctx, logger, err)
-
-		return resp, err
-	}
-}
-
-func addHeadersForResourceExhausted(ctx context.Context, logger log.Logger, err error) {
-	var reErr *serviceerror.ResourceExhausted
-	if errors.As(err, &reErr) {
-		headerErr := grpc.SetHeader(ctx, metadata.Pairs(
-			ResourceExhaustedCauseHeader, reErr.Cause.String(),
-			ResourceExhaustedScopeHeader, reErr.Scope.String(),
-		))
-		if headerErr != nil {
-			logger.Error("Failed to add Resource-Exhausted headers to response", tag.Error(headerErr))
-		}
-	}
-}

common/rpc/interceptor/frontend_service_error.go

@@ -0,0 +1,60 @@
+package interceptor
+
+import (
+	"context"
+
+	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/server/common/api"
+	"go.temporal.io/server/common/log"
+	"go.temporal.io/server/common/log/tag"
+	serviceerrors "go.temporal.io/server/common/serviceerror"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	// ResourceExhaustedCauseHeader is added to rpc response if request returns ResourceExhausted error.
+	ResourceExhaustedCauseHeader = "X-Resource-Exhausted-Cause"
+
+	// ResourceExhaustedScopeHeader is added to rpc response if request returns ResourceExhausted error.
+	ResourceExhaustedScopeHeader = "X-Resource-Exhausted-Scope"
+)
+
+// NewFrontendServiceErrorInterceptor returns a gRPC interceptor that has two responsibilities:
+//  1. Mask certain internal service error details.
+//  2. Propagate resource exhaustion details via gRPC headers.
+func NewFrontendServiceErrorInterceptor(
+	logger log.Logger,
+) grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		resp, err := handler(ctx, req)
+		if err == nil {
+			return resp, nil
+		}
+
+		switch serviceErr := err.(type) {
+		case *serviceerrors.ShardOwnershipLost:
+			err = serviceerror.NewUnavailable("shard unavailable, please backoff and retry")
+		case *serviceerror.DataLoss:
+			err = serviceerror.NewUnavailable("internal history service error")
+		case *serviceerror.ResourceExhausted:
+			if headerErr := grpc.SetHeader(ctx, metadata.Pairs(
+				ResourceExhaustedCauseHeader, serviceErr.Cause.String(),
+				ResourceExhaustedScopeHeader, serviceErr.Scope.String(),
+			)); headerErr != nil {
+				// So while this is *not* a user-facing error or problem in itself,
+				// it indicates that there might be larger connection issues at play.
+				logger.Error("Failed to add Resource-Exhausted headers to response",
+					tag.Operation(api.MethodName(info.FullMethod)),
+					tag.Error(headerErr))
+			}
+		}
+
+		return resp, err
+	}
+}

common/rpc/interceptor/frontend_service_error_test.go

@@ -0,0 +1,119 @@
+package interceptor
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	enumspb "go.temporal.io/api/enums/v1"
+	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/server/common/log/tag"
+	serviceerrors "go.temporal.io/server/common/serviceerror"
+	"go.temporal.io/server/common/testing/rpctest"
+	"go.temporal.io/server/common/testing/testlogger"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestFrontendServiceErrorInterceptor(t *testing.T) {
+	tests := []struct {
+		name            string
+		handlerErr      error
+		configureStream func(s *rpctest.MockServerTransportStream)
+		verifyFn        func(t *testing.T, err error, stream *rpctest.MockServerTransportStream)
+		expectLogErr    string
+	}{
+		{
+			name:       "Passthrough",
+			handlerErr: nil,
+			verifyFn: func(t *testing.T, err error, _ *rpctest.MockServerTransportStream) {
+				require.NoError(t, err)
+			},
+		},
+		{
+			name:       "Mask ShardOwnershipLost",
+			handlerErr: serviceerrors.NewShardOwnershipLost("owner-host", "current-host"),
+			verifyFn: func(t *testing.T, err error, _ *rpctest.MockServerTransportStream) {
+				require.Error(t, err)
+
+				var unavail *serviceerror.Unavailable
+				require.ErrorAs(t, err, &unavail)
+				assert.Contains(t, unavail.Error(), "shard unavailable")
+			},
+		},
+		{
+			name:       "Mask DataLoss",
+			handlerErr: serviceerror.NewDataLoss("..."),
+			verifyFn: func(t *testing.T, err error, _ *rpctest.MockServerTransportStream) {
+				require.Error(t, err)
+
+				var unavail *serviceerror.Unavailable
+				require.ErrorAs(t, err, &unavail)
+				assert.Equal(t, "internal history service error", unavail.Error())
+			},
+		},
+		{
+			name: "Set ResourceExhaustedHeaders",
+			handlerErr: &serviceerror.ResourceExhausted{
+				Cause: enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT,
+				Scope: enumspb.RESOURCE_EXHAUSTED_SCOPE_SYSTEM,
+			},
+			verifyFn: func(t *testing.T, err error, s *rpctest.MockServerTransportStream) {
+				require.Error(t, err)
+
+				hdr := s.CapturedHeaders()
+				require.NotNil(t, hdr)
+				assert.Equal(t, []string{
+					enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT.String()},
+					hdr.Get(ResourceExhaustedCauseHeader))
+				assert.Equal(t, []string{
+					enumspb.RESOURCE_EXHAUSTED_SCOPE_SYSTEM.String()},
+					hdr.Get(ResourceExhaustedScopeHeader))
+			},
+		},
+		{
+			name:       "Set ResourceExhaustedHeaders Failure",
+			handlerErr: serviceerror.NewResourceExhausted(enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT, "rate limit exceeded"),
+			configureStream: func(s *rpctest.MockServerTransportStream) {
+				s.SetHeaderFunc = func(md metadata.MD) error { return errors.New("injected header failure") }
+			},
+			expectLogErr: "Failed to add Resource-Exhausted headers to response",
+			verifyFn: func(t *testing.T, err error, _ *rpctest.MockServerTransportStream) {
+				require.Error(t, err)
+
+				var re *serviceerror.ResourceExhausted
+				require.ErrorAs(t, err, &re)
+				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT, re.Cause)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			method := "/test/method"
+
+			tl := testlogger.NewTestLogger(t, testlogger.FailOnAnyUnexpectedError)
+			if tc.expectLogErr != "" {
+				tl.Expect(testlogger.Error, tc.expectLogErr, tag.Operation("method"))
+			}
+
+			stream := rpctest.NewMockServerTransportStream(method)
+			if tc.configureStream != nil {
+				tc.configureStream(stream)
+			}
+			ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream)
+
+			var interceptorFn grpc.UnaryServerInterceptor
+			interceptorFn = NewFrontendServiceErrorInterceptor(tl)
+			info := &grpc.UnaryServerInfo{FullMethod: method}
+			_, err := interceptorFn(ctx, nil, info,
+				func(_ context.Context, _ any) (any, error) {
+					return nil, tc.handlerErr
+				})
+
+			tc.verifyFn(t, err, stream)
+		})
+	}
+}