Updates to partition config middleware (#5334)

Author: Shaddoll

common/rpc.go

@@ -56,8 +56,8 @@ const (
 	// IsolationGroupHeaderName refers to the name of the header that contains the isolation group of the client
 	IsolationGroupHeaderName = "cadence-worker-isolation-group"
 
-	// ClientZoneHeaderName refers to the name of the header that contaains the zone which the client request is from
-	ClientZoneHeaderName = "cadence-client-zone"
+	// ClientIsolationGroupHeaderName refers to the name of the header that contains the isolation group which the client request is from
+	ClientIsolationGroupHeaderName = "cadence-client-isolation-group"
 )
 
 type (

common/rpc/middleware.go

@@ -224,12 +224,12 @@ func (m *ForwardPartitionConfigMiddleware) Call(ctx context.Context, request *tr
 	return out.Call(ctx, request)
 }
 
-// ZonalPartitionConfigMiddleware stores the partition config and isolation group of the request into the context
-// It reads a header from client request which indicates the zone which the request is from and uses it as the isolation group
-type ZonalPartitionConfigMiddleware struct{}
+// ClientPartitionConfigMiddleware stores the partition config and isolation group of the request into the context
+// It reads a header from client request and uses it as the isolation group
+type ClientPartitionConfigMiddleware struct{}
 
-func (m *ZonalPartitionConfigMiddleware) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, h transport.UnaryHandler) error {
-	zone, _ := req.Headers.Get(common.ClientZoneHeaderName)
+func (m *ClientPartitionConfigMiddleware) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, h transport.UnaryHandler) error {
+	zone, _ := req.Headers.Get(common.ClientIsolationGroupHeaderName)
 	if zone != "" {
 		ctx = partition.ContextWithConfig(ctx, map[string]string{
 			partition.IsolationGroupKey: zone,

common/rpc/middleware_test.go

@@ -304,20 +304,20 @@ func TestForwardPartitionConfigMiddleware(t *testing.T) {
 	})
 }
 
-func TestZonalPartitionConfigMiddleware(t *testing.T) {
+func TestClientPartitionConfigMiddleware(t *testing.T) {
 	t.Run("it sets the partition config", func(t *testing.T) {
-		m := &ZonalPartitionConfigMiddleware{}
+		m := &ClientPartitionConfigMiddleware{}
 		h := &fakeHandler{}
 		headers := transport.NewHeaders().
-			With(common.ClientZoneHeaderName, "dca1")
+			With(common.ClientIsolationGroupHeaderName, "dca1")
 		err := m.Handle(context.Background(), &transport.Request{Headers: headers}, nil, h)
 		assert.NoError(t, err)
 		assert.Equal(t, map[string]string{partition.IsolationGroupKey: "dca1"}, partition.ConfigFromContext(h.ctx))
 		assert.Equal(t, "dca1", partition.IsolationGroupFromContext(h.ctx))
 	})
 
 	t.Run("noop when header is empty", func(t *testing.T) {
-		m := &ZonalPartitionConfigMiddleware{}
+		m := &ClientPartitionConfigMiddleware{}
 		h := &fakeHandler{}
 		headers := transport.NewHeaders()
 		ctx := context.Background()

common/rpc/params.go

@@ -132,12 +132,13 @@ func NewParams(serviceName string, config *config.Config, dc *dynamicconfig.Coll
 		InboundTLS:  inboundTLS,
 		OutboundTLS: outboundTLS,
 		InboundMiddleware: yarpc.InboundMiddleware{
-			Unary: &InboundMetricsMiddleware{},
+			// order matters: ForwardPartitionConfigMiddleware must be applied after ClientPartitionConfigMiddleware
+			Unary: yarpc.UnaryInboundMiddleware(&InboundMetricsMiddleware{}, &ClientPartitionConfigMiddleware{}, &ForwardPartitionConfigMiddleware{}),
 		},
 		OutboundMiddleware: yarpc.OutboundMiddleware{
-			Unary: &HeaderForwardingMiddleware{
+			Unary: yarpc.UnaryOutboundMiddleware(&HeaderForwardingMiddleware{
 				Rules: forwardingRules,
-			},
+			}, &ForwardPartitionConfigMiddleware{}),
 		},
 	}, nil
 }