more changes

client/clientfactory.go

@@ -42,7 +42,6 @@ import (
 	"github.com/uber/cadence/client/wrappers/metered"
 	"github.com/uber/cadence/client/wrappers/thrift"
 	timeoutwrapper "github.com/uber/cadence/client/wrappers/timeout"
-	"github.com/uber/cadence/common"
 	"github.com/uber/cadence/common/dynamicconfig"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/membership"
@@ -68,7 +67,7 @@ type (
 	DomainIDToNameFunc func(string) (string, error)
 
 	rpcClientFactory struct {
-		rpcFactory            common.RPCFactory
+		rpcFactory            rpc.Factory
 		resolver              membership.Resolver
 		metricsClient         metrics.Client
 		dynConfig             *dynamicconfig.Collection
@@ -79,7 +78,7 @@ type (
 
 // NewRPCClientFactory creates an instance of client factory that knows how to dispatch RPC calls.
 func NewRPCClientFactory(
-	rpcFactory common.RPCFactory,
+	rpcFactory rpc.Factory,
 	resolver membership.Resolver,
 	metricsClient metrics.Client,
 	dc *dynamicconfig.Collection,

common/dynamicconfig/constants.go

@@ -1588,6 +1588,12 @@ const (
 	// Default value: false
 	// Allowed filters: N/A
 	EnableSQLAsyncTransaction
+	// EnableConnectionRetainingDirectChooser is the key for enabling connection retaining direct yarpc chooser
+	// KeyName: system.enableConnectionRetainingDirectChooser
+	// Value type: Bool
+	// Default value: false
+	// Allowed filters: N/A
+	EnableConnectionRetainingDirectChooser
 
 	// key for frontend
 
@@ -3950,6 +3956,11 @@ var BoolKeys = map[BoolKey]DynamicBool{
 		Description:  "EnableSQLAsyncTransaction is the key for enabling async transaction",
 		DefaultValue: false,
 	},
+	EnableConnectionRetainingDirectChooser: {
+		KeyName:      "system.enableConnectionRetainingDirectChooser",
+		Description:  "EnableConnectionRetainingDirectChooser is the key for enabling connection retaining direct chooser",
+		DefaultValue: false,
+	},
 	EnableClientVersionCheck: {
 		KeyName:      "frontend.enableClientVersionCheck",
 		Description:  "EnableClientVersionCheck is enables client version check for frontend",

common/rpc.go → common/headers.go

@@ -18,14 +18,8 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-//go:generate mockgen -package $GOPACKAGE -source $GOFILE -destination rpc_mock.go -self_package github.com/uber/cadence/common
-
 package common
 
-import (
-	"go.uber.org/yarpc"
-)
-
 const (
 	// LibraryVersionHeaderName refers to the name of the
 	// tchannel / http header that contains the client
@@ -61,13 +55,3 @@ const (
 	// 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 (
-	// RPCFactory Creates a dispatcher that knows how to transport requests.
-	RPCFactory interface {
-		GetDispatcher() *yarpc.Dispatcher
-		GetMaxMessageSize() int
-		Start()
-		Stop()
-	}
-)

common/log/tag/values.go

@@ -138,6 +138,7 @@ var (
 	ComponentMapQ                       = component("mapq")
 	ComponentMapQTree                   = component("mapq-tree")
 	ComponentMapQTreeNode               = component("mapq-tree-node")
+	ComponentRPCFactory                 = component("rpc-factory")
 )
 
 // Pre-defined values for TagSysLifecycle

common/membership/resolver.go

@@ -48,6 +48,7 @@ type (
 	// Resolver provides membership information for all cadence services.
 	Resolver interface {
 		common.Daemon
+
 		// WhoAmI returns self host details.
 		// To be consistent with peer provider, it is advised to use peer provider
 		// to return this information

common/resource/params.go

@@ -44,6 +44,7 @@ import (
 	"github.com/uber/cadence/common/metrics"
 	"github.com/uber/cadence/common/partition"
 	"github.com/uber/cadence/common/pinot"
+	"github.com/uber/cadence/common/rpc"
 )
 
 type (
@@ -58,7 +59,7 @@ type (
 
 		MetricScope        tally.Scope
 		MembershipResolver membership.Resolver
-		RPCFactory         common.RPCFactory
+		RPCFactory         rpc.Factory
 		PProfInitializer   common.PProfInitializer
 		PersistenceConfig  config.Persistence
 		ClusterMetadata    cluster.Metadata

common/resource/resource.go

@@ -49,7 +49,7 @@ import (
 	"github.com/uber/cadence/common/partition"
 	"github.com/uber/cadence/common/persistence"
 	persistenceClient "github.com/uber/cadence/common/persistence/client"
-	"github.com/uber/cadence/common/quotas/global/rpc"
+	qrpc "github.com/uber/cadence/common/quotas/global/rpc"
 	"github.com/uber/cadence/common/service"
 )
 
@@ -96,7 +96,7 @@ type (
 		GetMatchingClient() matching.Client
 		GetHistoryRawClient() history.Client
 		GetHistoryClient() history.Client
-		GetRatelimiterAggregatorsClient() rpc.Client
+		GetRatelimiterAggregatorsClient() qrpc.Client
 		GetRemoteAdminClient(cluster string) admin.Client
 		GetRemoteFrontendClient(cluster string) frontend.Client
 		GetClientBean() client.Bean

common/resource/resourceImpl.go

@@ -58,8 +58,9 @@ import (
 	"github.com/uber/cadence/common/partition"
 	"github.com/uber/cadence/common/persistence"
 	persistenceClient "github.com/uber/cadence/common/persistence/client"
-	"github.com/uber/cadence/common/quotas/global/rpc"
+	qrpc "github.com/uber/cadence/common/quotas/global/rpc"
 	"github.com/uber/cadence/common/quotas/permember"
+	"github.com/uber/cadence/common/rpc"
 	"github.com/uber/cadence/common/service"
 )
 
@@ -141,15 +142,15 @@ type (
 
 		pprofInitializer       common.PProfInitializer
 		runtimeMetricsReporter *metrics.RuntimeMetricsReporter
-		rpcFactory             common.RPCFactory
+		rpcFactory             rpc.Factory
 
 		isolationGroups           isolationgroup.State
 		isolationGroupConfigStore configstore.Client
 		partitioner               partition.Partitioner
 
 		asyncWorkflowQueueProvider queue.Provider
 
-		ratelimiterAggregatorClient rpc.Client
+		ratelimiterAggregatorClient qrpc.Client
 	}
 )
 
@@ -304,7 +305,7 @@ func New(
 	}
 	partitioner := ensurePartitionerOrDefault(params, isolationGroupState)
 
-	ratelimiterAggs := rpc.New(
+	ratelimiterAggs := qrpc.New(
 		historyRawClient, // no retries, will retry internally if needed
 		clientBean.GetHistoryPeers(),
 		logger,
@@ -399,7 +400,7 @@ func (h *Impl) Start() {
 		h.logger.WithTags(tag.Error(err)).Fatal("fail to start PProf")
 	}
 
-	h.rpcFactory.Start()
+	h.rpcFactory.Start(h.membershipResolver)
 
 	if err := h.dispatcher.Start(); err != nil {
 		h.logger.WithTags(tag.Error(err)).Fatal("fail to start dispatcher")
@@ -425,7 +426,6 @@ func (h *Impl) Start() {
 
 // Stop stops all resources
 func (h *Impl) Stop() {
-
 	if !atomic.CompareAndSwapInt32(
 		&h.status,
 		common.DaemonStatusStarted,
@@ -560,7 +560,7 @@ func (h *Impl) GetHistoryClient() history.Client {
 	return h.historyClient
 }
 
-func (h *Impl) GetRatelimiterAggregatorsClient() rpc.Client {
+func (h *Impl) GetRatelimiterAggregatorsClient() qrpc.Client {
 	return h.ratelimiterAggregatorClient
 }
 

common/rpc/direct_peer_chooser.go

@@ -0,0 +1,110 @@
+// Copyright (c) 2021 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package rpc
+
+import (
+	"context"
+	"sync"
+
+	"go.uber.org/yarpc/api/peer"
+	"go.uber.org/yarpc/api/transport"
+	"go.uber.org/yarpc/peer/direct"
+	"go.uber.org/yarpc/yarpcerrors"
+
+	"github.com/uber/cadence/common/dynamicconfig"
+	"github.com/uber/cadence/common/log"
+	"github.com/uber/cadence/common/membership"
+)
+
+// directPeerChooser is a peer.Chooser that chooses a peer based on the shard key.
+// Peers are managed by the peerList and peers are reused across multiple requests.
+type directPeerChooser struct {
+	serviceName          string
+	logger               log.Logger
+	t                    peer.Transport
+	enableConnRetainMode dynamicconfig.BoolPropertyFn
+	legacyChooser        peer.Chooser
+	mu                   sync.RWMutex
+}
+
+func newDirectChooser(serviceName string, t peer.Transport, logger log.Logger, enableConnRetainMode dynamicconfig.BoolPropertyFn) *directPeerChooser {
+	return &directPeerChooser{
+		serviceName:          serviceName,
+		logger:               logger,
+		t:                    t,
+		enableConnRetainMode: enableConnRetainMode,
+	}
+}
+
+// Start statisfies the peer.Chooser interface.
+func (*directPeerChooser) Start() error {
+	return nil // no-op
+}
+
+// Stop statisfies the peer.Chooser interface.
+func (*directPeerChooser) Stop() error {
+	return nil // no-op
+}
+
+// IsRunning statisfies the peer.Chooser interface.
+func (*directPeerChooser) IsRunning() bool {
+	return true // no-op
+}
+
+// Choose returns an existing peer for the shard key.
+func (g *directPeerChooser) Choose(ctx context.Context, req *transport.Request) (peer peer.Peer, onFinish func(error), err error) {
+	if g.enableConnRetainMode != nil && !g.enableConnRetainMode() {
+		return g.chooseFromLegacyDirectPeerChooser(ctx, req)
+	}
+
+	if req.ShardKey == "" {
+		return nil, nil, yarpcerrors.InvalidArgumentErrorf("chooser requires ShardKey to be non-empty")
+	}
+
+	// TODO: implement connection retain mode
+	return nil, nil, yarpcerrors.UnimplementedErrorf("direct peer chooser conn retain mode unimplemented")
+}
+
+func (g *directPeerChooser) UpdatePeers([]membership.HostInfo) {
+	// TODO: implement
+}
+
+func (g *directPeerChooser) chooseFromLegacyDirectPeerChooser(ctx context.Context, req *transport.Request) (peer.Peer, func(error), error) {
+	g.mu.RLock()
+
+	if g.legacyChooser != nil {
+		g.mu.RUnlock()
+		return g.legacyChooser.Choose(ctx, req)
+	}
+
+	g.mu.RUnlock()
+
+	var err error
+	g.mu.Lock()
+	g.legacyChooser, err = direct.New(direct.Configuration{}, g.t)
+	g.mu.Unlock()
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return g.legacyChooser.Choose(ctx, req)
+}