feat(frontend): support autoforwarding through headers

Signed-off-by: Shijie Sheng <liouvetren@gmail.com>

Author: shijiesheng

common/headers.go

@@ -59,4 +59,7 @@ const (
 
 	// CallerTypeHeaderName refers to the name of the header that contains the caller type (CLI, UI, SDK, internal, etc.)
 	CallerTypeHeaderName = types.CallerTypeHeaderName
+
+	// QueryConsistencyLevelHeaderName refers to the name of the header that contains the query consistency level (EVENTUAL, STRONG)
+	QueryConsistencyLevelHeaderName = "cadence-query-consistency-level"
 )

common/rpc/middleware.go

@@ -23,10 +23,12 @@ package rpc
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 
 	"go.uber.org/cadence/worker"
 	"go.uber.org/yarpc"
+	"go.uber.org/yarpc/api/middleware"
 	"go.uber.org/yarpc/api/transport"
 
 	"github.com/uber/cadence/common"
@@ -36,6 +38,8 @@ import (
 	"github.com/uber/cadence/common/types"
 )
 
+var _ middleware.UnaryInbound = &QueryConsistencyLevelInboundMiddleware{}
+
 type authOutboundMiddleware struct {
 	authProvider worker.AuthorizationProvider
 }
@@ -247,3 +251,18 @@ func (m *ClientPartitionConfigMiddleware) Handle(ctx context.Context, req *trans
 	}
 	return h.Handle(ctx, req, resw)
 }
+
+type QueryConsistencyLevelInboundMiddleware struct{}
+
+func (m *QueryConsistencyLevelInboundMiddleware) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, h transport.UnaryHandler) error {
+	queryConsistencyLevel, ok := req.Headers.Get(common.QueryConsistencyLevelHeaderName)
+	if ok {
+		var level types.QueryConsistencyLevel
+		err := level.UnmarshalText([]byte(queryConsistencyLevel))
+		if err != nil {
+			return fmt.Errorf("failed to parse query consistency level: %w", err)
+		}
+		ctx = context.WithValue(ctx, common.QueryConsistencyLevelHeaderName, level)
+	}
+	return h.Handle(ctx, req, resw)
+}

common/rpc/middleware_test.go

@@ -381,6 +381,99 @@ func TestCallerInfoMiddleware(t *testing.T) {
 	})
 }
 
+func TestQueryConsistencyLevelInboundMiddleware(t *testing.T) {
+	tests := []struct {
+		name          string
+		headerValue   string
+		hasHeader     bool
+		expectedLevel types.QueryConsistencyLevel
+		expectError   bool
+	}{
+		{
+			name:          "parses EVENTUAL consistency level",
+			headerValue:   "EVENTUAL",
+			hasHeader:     true,
+			expectedLevel: types.QueryConsistencyLevelEventual,
+			expectError:   false,
+		},
+		{
+			name:          "parses STRONG consistency level",
+			headerValue:   "STRONG",
+			hasHeader:     true,
+			expectedLevel: types.QueryConsistencyLevelStrong,
+			expectError:   false,
+		},
+		{
+			name:          "parses lowercase eventual",
+			headerValue:   "eventual",
+			hasHeader:     true,
+			expectedLevel: types.QueryConsistencyLevelEventual,
+			expectError:   false,
+		},
+		{
+			name:          "parses lowercase strong",
+			headerValue:   "strong",
+			hasHeader:     true,
+			expectedLevel: types.QueryConsistencyLevelStrong,
+			expectError:   false,
+		},
+		{
+			name:          "parses numeric value 0",
+			headerValue:   "0",
+			hasHeader:     true,
+			expectedLevel: types.QueryConsistencyLevelEventual,
+			expectError:   false,
+		},
+		{
+			name:          "parses numeric value 1",
+			headerValue:   "1",
+			hasHeader:     true,
+			expectedLevel: types.QueryConsistencyLevelStrong,
+			expectError:   false,
+		},
+		{
+			name:        "returns error for invalid value",
+			headerValue: "INVALID",
+			hasHeader:   true,
+			expectError: true,
+		},
+		{
+			name:        "no header - noop",
+			hasHeader:   false,
+			expectError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			m := &QueryConsistencyLevelInboundMiddleware{}
+			h := &fakeHandler{}
+			headers := transport.NewHeaders()
+			if tt.hasHeader {
+				headers = headers.With(common.QueryConsistencyLevelHeaderName, tt.headerValue)
+			}
+
+			err := m.Handle(context.Background(), &transport.Request{Headers: headers}, nil, h)
+
+			if tt.expectError {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), "failed to parse query consistency level")
+				return
+			}
+
+			assert.NoError(t, err)
+			if tt.hasHeader {
+				level, ok := h.ctx.Value(common.QueryConsistencyLevelHeaderName).(types.QueryConsistencyLevel)
+				assert.True(t, ok, "expected query consistency level to be set in context")
+				assert.Equal(t, tt.expectedLevel, level)
+			} else {
+				// When no header is present, context value should not be set
+				assert.Nil(t, h.ctx.Value(common.QueryConsistencyLevelHeaderName))
+			}
+		})
+	}
+}
+
 type fakeHandler struct {
 	ctx context.Context
 }

common/rpc/params.go

@@ -170,7 +170,7 @@ func NewParams(serviceName string, config *config.Config, dc *dynamicconfig.Coll
 		OutboundTLS:      outboundTLS,
 		InboundMiddleware: yarpc.InboundMiddleware{
 			// order matters: ForwardPartitionConfigMiddleware must be applied after ClientPartitionConfigMiddleware
-			Unary: yarpc.UnaryInboundMiddleware(&InboundMetricsMiddleware{}, &CallerInfoMiddleware{}, &ClientPartitionConfigMiddleware{}, &ForwardPartitionConfigMiddleware{}),
+			Unary: yarpc.UnaryInboundMiddleware(&InboundMetricsMiddleware{}, &CallerInfoMiddleware{}, &ClientPartitionConfigMiddleware{}, &ForwardPartitionConfigMiddleware{}, &QueryConsistencyLevelInboundMiddleware{}),
 		},
 		OutboundMiddleware: yarpc.OutboundMiddleware{
 			Unary: yarpc.UnaryOutboundMiddleware(&HeaderForwardingMiddleware{

host/onebox.go

@@ -1159,7 +1159,7 @@ func (c *cadenceImpl) newRPCFactory(serviceName string, host membership.HostInfo
 		TChannelAddress: tchannelAddress,
 		GRPCAddress:     grpcAddress,
 		InboundMiddleware: yarpc.InboundMiddleware{
-			Unary: yarpc.UnaryInboundMiddleware(&versionMiddleware{}, &rpc.ClientPartitionConfigMiddleware{}, &rpc.ForwardPartitionConfigMiddleware{}),
+			Unary: yarpc.UnaryInboundMiddleware(&versionMiddleware{}, &rpc.ClientPartitionConfigMiddleware{}, &rpc.ForwardPartitionConfigMiddleware{}, &rpc.QueryConsistencyLevelInboundMiddleware{}),
 		},
 		OutboundMiddleware: yarpc.OutboundMiddleware{
 			Unary: &rpc.ForwardPartitionConfigMiddleware{},

service/frontend/templates/clusterredirection.tmpl

@@ -70,7 +70,7 @@ func (handler *clusterRedirectionHandler) {{$method.Declaration}} {
 	var (
 		apiName = "{{$method.Name}}"
 		cluster string
-		requestedConsistencyLevel types.QueryConsistencyLevel = types.QueryConsistencyLevelEventual
+		requestedConsistencyLevel types.QueryConsistencyLevel = getRequestedConsistencyLevelFromContext(ctx)
 	)
 
 	{{- if has $method.Name $readAPIsWithStrongConsistency}}