GetWorkflowExecutionHistory long poll soft timeout (#8238)

## What changed?

Added a "soft timeout" (language used by Workflow Update) to
`GetWorkflowExecutionHistory` long polls.

## Why?

We don't want to terminate the long poll connection but instead keep it
alive by sending a response back just before the timeout. The idea is
that this will prevent connections from opening/terminating repeatedly
(ie connection churn).

## How did you test it?
- [ ] built
- [x] run locally and tested manually
- [x] covered by existing tests
- [ ] added new unit test(s)
- [ ] added new functional test(s)

## Potential risks

I was only able to verify this manually by forcing a timeout in the
server and verifying that instead of a deadline exceeded I saw a result.
I'll assume this will work since the existing code already tried doing
just exactly that, but it didn't do it well.

.github/.golangci.yml

@@ -20,6 +20,7 @@ linters:
     staticcheck:
       checks:
         - "all"
+        - "-ST1000" # disable: package comment is missing
     godox:
       keywords:
         - FIXME

common/contextutil/deadline.go

@@ -0,0 +1,37 @@
+package contextutil
+
+import (
+	"context"
+	"time"
+)
+
+var noop = func() {}
+
+// WithDeadlineBuffer creates a child context with desired timeout.
+// If buffer is non-zero, then child context timeout will be
+// the minOf(parentCtx.Deadline()-buffer, maxTimeout). Use this
+// method to create child context when childContext cannot use
+// all of parent's deadline but instead there is a need to leave
+// some time for parent to do some post-work
+func WithDeadlineBuffer(
+	parent context.Context,
+	timeout time.Duration,
+	buffer time.Duration,
+) (context.Context, context.CancelFunc) {
+	if parent.Err() != nil {
+		return parent, noop
+	}
+
+	deadline, hasDeadline := parent.Deadline()
+
+	if !hasDeadline {
+		return context.WithTimeout(parent, timeout)
+	}
+
+	remaining := time.Until(deadline) - buffer
+	if remaining < timeout {
+		// Cap the timeout to the remaining time minus buffer.
+		timeout = max(0, remaining)
+	}
+	return context.WithTimeout(parent, timeout)
+}

common/contextutil/deadline_test.go

@@ -0,0 +1,59 @@
+package contextutil
+
+import (
+	"context"
+	"math"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+const testTolerance = 5 * time.Second
+
+func TestWithDeadlineBuffer(t *testing.T) {
+	const timeout = 10 * time.Minute
+	const buffer = 1 * time.Minute
+	start := time.Now()
+
+	t.Run("parent is cancelled", func(t *testing.T) {
+		parent, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		child, _ := WithDeadlineBuffer(parent, timeout, buffer)
+		require.Equal(t, parent, child)
+	})
+
+	t.Run("parent has no deadline", func(t *testing.T) {
+		parent := context.Background()
+
+		t.Run("timeout specified", func(t *testing.T) {
+			child, _ := WithDeadlineBuffer(parent, timeout, 0)
+			dl, _ := child.Deadline()
+			require.WithinDuration(t, start.Add(timeout), dl, testTolerance)
+		})
+	})
+
+	t.Run("parent has deadline", func(t *testing.T) {
+		parent, parentCancel := context.WithTimeout(context.Background(), timeout)
+		defer parentCancel()
+		parentDeadline, _ := parent.Deadline()
+
+		t.Run("enough buffer left", func(t *testing.T) {
+			child, _ := WithDeadlineBuffer(parent, math.MaxInt, buffer)
+			dl, _ := child.Deadline()
+			require.WithinDuration(t, parentDeadline.Add(-buffer), dl, testTolerance)
+		})
+
+		t.Run("no buffer left", func(t *testing.T) {
+			child, _ := WithDeadlineBuffer(parent, math.MaxInt, math.MaxInt)
+			require.Equal(t, child.Err(), context.DeadlineExceeded)
+		})
+
+		t.Run("enough buffer left but less than max timeout", func(t *testing.T) {
+			child, _ := WithDeadlineBuffer(parent, timeout/2, buffer)
+			dl, _ := child.Deadline()
+			require.WithinDuration(t, parentDeadline.Add(-timeout/2), dl, testTolerance)
+		})
+	})
+}

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,13 +3,10 @@ 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"
@@ -17,7 +14,6 @@ import (
 	"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,118 @@
+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 = 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)
+		})
+	}
+}

common/testing/rpctest/transport_stream.go

@@ -0,0 +1,98 @@
+package rpctest
+
+import (
+	"sync"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+var _ grpc.ServerTransportStream = (*MockServerTransportStream)(nil)
+
+// MockServerTransportStream is a reusable test double that mimics gRPC's
+// internal ServerTransportStream.
+type MockServerTransportStream struct {
+	mu       sync.Mutex
+	method   string
+	headers  metadata.MD
+	trailers []*metadata.MD
+
+	// Optional hook overrides. If non-nil they are invoked instead of the
+	// default header / send header logic.
+	SetHeaderFunc  func(metadata.MD) error
+	SendHeaderFunc func(metadata.MD) error
+	TrailerFunc    func(metadata.MD) error
+}
+
+func NewMockServerTransportStream(methodName string) *MockServerTransportStream {
+	return &MockServerTransportStream{
+		method: methodName,
+	}
+}
+
+func (s *MockServerTransportStream) Method() string {
+	return s.method
+}
+
+func (s *MockServerTransportStream) SetHeader(md metadata.MD) error {
+	if s.SetHeaderFunc != nil {
+		return s.SetHeaderFunc(md)
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.headers == nil {
+		s.headers = metadata.New(nil)
+	}
+	for k, v := range md {
+		s.headers[k] = append(s.headers[k], v...)
+	}
+	return nil
+}
+
+func (s *MockServerTransportStream) SendHeader(md metadata.MD) error {
+	if s.SendHeaderFunc != nil {
+		return s.SendHeaderFunc(md)
+	}
+	return s.SetHeader(md)
+}
+
+func (s *MockServerTransportStream) SetTrailer(md metadata.MD) error {
+	if s.TrailerFunc != nil {
+		return s.TrailerFunc(md)
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	cp := md.Copy()
+	s.trailers = append(s.trailers, &cp)
+	return nil
+}
+
+func (s *MockServerTransportStream) CapturedHeaders() metadata.MD {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.headers == nil {
+		return nil
+	}
+	out := metadata.New(nil)
+	for k, v := range s.headers {
+		out[k] = append([]string(nil), v...)
+	}
+	return out
+}
+
+func (s *MockServerTransportStream) CapturedTrailers() []metadata.MD {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	out := make([]metadata.MD, 0, len(s.trailers))
+	for _, t := range s.trailers {
+		if t == nil {
+			continue
+		}
+		cp := metadata.New(nil)
+		for k, v := range *t {
+			cp[k] = append([]string(nil), v...)
+		}
+		out = append(out, cp)
+	}
+	return out
+}

service/frontend/fx.go

@@ -30,7 +30,6 @@ import (
 	"go.temporal.io/server/common/quotas/calculator"
 	"go.temporal.io/server/common/resolver"
 	"go.temporal.io/server/common/resource"
-	"go.temporal.io/server/common/rpc"
 	"go.temporal.io/server/common/rpc/encryption"
 	"go.temporal.io/server/common/rpc/interceptor"
 	"go.temporal.io/server/common/sdk"
@@ -239,7 +238,7 @@ func GrpcServerOptionsProvider(
 		// Service Error Interceptor should be the next most outer interceptor on error handling
 		maskInternalErrorDetailsInterceptor.Intercept,
 		interceptor.ServiceErrorInterceptor,
-		rpc.NewFrontendServiceErrorInterceptor(logger),
+		interceptor.NewFrontendServiceErrorInterceptor(logger),
 		namespaceValidatorInterceptor.NamespaceValidateIntercept,
 		namespaceLogInterceptor.Intercept, // TODO: Deprecate this with a outer custom interceptor
 		metrics.NewServerMetricsContextInjectorInterceptor(),

service/frontend/fx_test.go

@@ -18,7 +18,6 @@ import (
 	"go.temporal.io/server/common/metrics/metricstest"
 	"go.temporal.io/server/common/namespace"
 	"go.temporal.io/server/common/primitives"
-	"go.temporal.io/server/common/rpc"
 	"go.temporal.io/server/common/rpc/interceptor"
 	"go.temporal.io/server/common/testing/nettest"
 	"go.uber.org/mock/gomock"
@@ -227,7 +226,7 @@ func TestRateLimitInterceptorProvider(t *testing.T) {
 			svc := &testSvc{}
 			server := grpc.NewServer(grpc.ChainUnaryInterceptor(
 				interceptor.ServiceErrorInterceptor,
-				rpc.NewFrontendServiceErrorInterceptor(log.NewTestLogger()),
+				interceptor.NewFrontendServiceErrorInterceptor(log.NewTestLogger()),
 				rateLimitInterceptor.Intercept,
 			))
 			workflowservice.RegisterWorkflowServiceServer(server, svc)
@@ -279,10 +278,10 @@ func TestRateLimitInterceptorProvider(t *testing.T) {
 				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT, resourceExhausted.Cause)
 				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_SCOPE_SYSTEM, resourceExhausted.Scope)
 
-				assert.Len(t, header.Get(rpc.ResourceExhaustedCauseHeader), 1)
-				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT.String(), header.Get(rpc.ResourceExhaustedCauseHeader)[0])
-				assert.Len(t, header.Get(rpc.ResourceExhaustedScopeHeader), 1)
-				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_SCOPE_SYSTEM.String(), header.Get(rpc.ResourceExhaustedScopeHeader)[0])
+				assert.Len(t, header.Get(interceptor.ResourceExhaustedCauseHeader), 1)
+				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT.String(), header.Get(interceptor.ResourceExhaustedCauseHeader)[0])
+				assert.Len(t, header.Get(interceptor.ResourceExhaustedScopeHeader), 1)
+				assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_SCOPE_SYSTEM.String(), header.Get(interceptor.ResourceExhaustedScopeHeader)[0])
 			} else {
 				assert.NoError(t, err)
 			}
@@ -583,7 +582,7 @@ func TestNamespaceRateLimitInterceptorProvider(t *testing.T) {
 			svc := &testSvc{}
 			server := grpc.NewServer(grpc.ChainUnaryInterceptor(
 				interceptor.ServiceErrorInterceptor,
-				rpc.NewFrontendServiceErrorInterceptor(log.NewTestLogger()),
+				interceptor.NewFrontendServiceErrorInterceptor(log.NewTestLogger()),
 				rateLimitInterceptor.Intercept,
 			))
 			workflowservice.RegisterWorkflowServiceServer(server, svc)
@@ -624,10 +623,10 @@ func TestNamespaceRateLimitInterceptorProvider(t *testing.T) {
 					assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT, resourceExhausted.Cause)
 					assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_SCOPE_NAMESPACE, resourceExhausted.Scope)
 
-					assert.Len(t, header.Get(rpc.ResourceExhaustedCauseHeader), 1)
-					assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT.String(), header.Get(rpc.ResourceExhaustedCauseHeader)[0])
-					assert.Len(t, header.Get(rpc.ResourceExhaustedScopeHeader), 1)
-					assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_SCOPE_NAMESPACE.String(), header.Get(rpc.ResourceExhaustedScopeHeader)[0])
+					assert.Len(t, header.Get(interceptor.ResourceExhaustedCauseHeader), 1)
+					assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT.String(), header.Get(interceptor.ResourceExhaustedCauseHeader)[0])
+					assert.Len(t, header.Get(interceptor.ResourceExhaustedScopeHeader), 1)
+					assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_SCOPE_NAMESPACE.String(), header.Get(interceptor.ResourceExhaustedScopeHeader)[0])
 				} else {
 					assert.NoError(t, err)
 				}
@@ -773,7 +772,7 @@ func TestNamespaceRateLimitMetrics(t *testing.T) {
 			svc := &testSvc{}
 			server := grpc.NewServer(grpc.ChainUnaryInterceptor(
 				interceptor.ServiceErrorInterceptor,
-				rpc.NewFrontendServiceErrorInterceptor(log.NewTestLogger()),
+				interceptor.NewFrontendServiceErrorInterceptor(log.NewTestLogger()),
 				rateLimitInterceptor.Intercept,
 			))
 			workflowservice.RegisterWorkflowServiceServer(server, svc)

service/history/api/get_workflow_util.go

@@ -12,6 +12,7 @@ import (
 	"go.temporal.io/server/api/historyservice/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	"go.temporal.io/server/common"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/locks"
 	"go.temporal.io/server/common/log/tag"
@@ -23,6 +24,8 @@ import (
 	historyi "go.temporal.io/server/service/history/interfaces"
 )
 
+const longPollSoftTimeout = time.Second
+
 func GetOrPollMutableState(
 	ctx context.Context,
 	shardContext historyi.ShardContext,
@@ -183,8 +186,12 @@ func GetOrPollMutableState(
 		if err != nil {
 			return nil, err
 		}
-		timer := time.NewTimer(shardContext.GetConfig().LongPollExpirationInterval(namespaceRegistry.Name().String()))
-		defer timer.Stop()
+
+		// Send back response just before caller context would time out.
+		longPollInterval := shardContext.GetConfig().LongPollExpirationInterval(namespaceRegistry.Name().String())
+		longPollCtx, cancel := contextutil.WithDeadlineBuffer(ctx, longPollInterval, longPollSoftTimeout)
+		defer cancel()
+
 		for {
 			select {
 			case event := <-channel:
@@ -239,10 +246,8 @@ func GetOrPollMutableState(
 				if expectedNextEventID < response.GetNextEventId() || response.GetWorkflowStatus() != enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING {
 					return response, nil
 				}
-			case <-timer.C:
+			case <-longPollCtx.Done():
 				return response, nil
-			case <-ctx.Done():
-				return nil, ctx.Err()
 			}
 		}
 	}

service/matching/matching_engine.go

@@ -34,6 +34,7 @@ import (
 	hlc "go.temporal.io/server/common/clock/hybrid_logical_clock"
 	"go.temporal.io/server/common/cluster"
 	"go.temporal.io/server/common/collection"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
@@ -2416,7 +2417,7 @@ func (e *matchingEngineImpl) ListNexusEndpoints(ctx context.Context, request *ma
 		request.LastKnownTableVersion = 0
 
 		var cancel context.CancelFunc
-		ctx, cancel = newChildContext(ctx, e.config.ListNexusEndpointsLongPollTimeout(), returnEmptyTaskTimeBudget)
+		ctx, cancel = contextutil.WithDeadlineBuffer(ctx, e.config.ListNexusEndpointsLongPollTimeout(), returnEmptyTaskTimeBudget)
 		defer cancel()
 	}
 
@@ -2548,7 +2549,7 @@ func (e *matchingEngineImpl) pollTask(
 	// reached, instead of emptyTask, context timeout error is returned to the frontend by the rpc stack,
 	// which counts against our SLO. By shortening the timeout by a very small amount, the emptyTask can be
 	// returned to the handler before a context timeout error is generated.
-	ctx, cancel := newChildContext(ctx, pm.LongPollExpirationInterval(), returnEmptyTaskTimeBudget)
+	ctx, cancel := contextutil.WithDeadlineBuffer(ctx, pm.LongPollExpirationInterval(), returnEmptyTaskTimeBudget)
 	defer cancel()
 
 	if pollerID, ok := ctx.Value(pollerIDKey).(string); ok && pollerID != "" {

service/matching/physical_task_queue_manager.go

@@ -20,6 +20,7 @@ import (
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/clock"
 	"go.temporal.io/server/common/cluster"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/debug"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
@@ -582,7 +583,7 @@ func (c *physicalTaskQueueManagerImpl) TrySyncMatch(ctx context.Context, task *i
 		return c.priMatcher.Offer(ctx, task)
 	}
 
-	childCtx, cancel := newChildContext(ctx, c.config.SyncMatchWaitDuration(), time.Second)
+	childCtx, cancel := contextutil.WithDeadlineBuffer(ctx, c.config.SyncMatchWaitDuration(), time.Second)
 	defer cancel()
 
 	return c.oldMatcher.Offer(childCtx, task)
@@ -693,31 +694,6 @@ func (c *physicalTaskQueueManagerImpl) ensureRegisteredInDeploymentVersion(
 	return nil
 }
 
-// newChildContext creates a child context with desired timeout.
-// if tailroom is non-zero, then child context timeout will be
-// the minOf(parentCtx.Deadline()-tailroom, timeout). Use this
-// method to create child context when childContext cannot use
-// all of parent's deadline but instead there is a need to leave
-// some time for parent to do some post-work
-func newChildContext(
-	parent context.Context,
-	timeout time.Duration,
-	tailroom time.Duration,
-) (context.Context, context.CancelFunc) {
-	if parent.Err() != nil {
-		return parent, func() {}
-	}
-	deadline, ok := parent.Deadline()
-	if !ok {
-		return context.WithTimeout(parent, timeout)
-	}
-	remaining := time.Until(deadline) - tailroom
-	if remaining < timeout {
-		timeout = max(0, remaining)
-	}
-	return context.WithTimeout(parent, timeout)
-}
-
 func (c *physicalTaskQueueManagerImpl) QueueKey() *PhysicalTaskQueueKey {
 	return c.queue
 }

service/matching/user_data_manager.go

@@ -15,6 +15,7 @@ import (
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/backoff"
 	"go.temporal.io/server/common/clock/hybrid_logical_clock"
+	"go.temporal.io/server/common/contextutil"
 	"go.temporal.io/server/common/future"
 	"go.temporal.io/server/common/goro"
 	"go.temporal.io/server/common/headers"
@@ -544,7 +545,7 @@ func (m *userDataManagerImpl) HandleGetUserDataRequest(
 
 	if req.WaitNewData {
 		var cancel context.CancelFunc
-		ctx, cancel = newChildContext(ctx, m.config.GetUserDataLongPollTimeout(), m.config.GetUserDataReturnBudget)
+		ctx, cancel = contextutil.WithDeadlineBuffer(ctx, m.config.GetUserDataLongPollTimeout(), m.config.GetUserDataReturnBudget)
 		defer cancel()
 	}