Introduce partition config provider to matching client

client/clientfactory.go

@@ -151,9 +151,10 @@ func (cf *rpcClientFactory) NewMatchingClientWithTimeout(
 
 	peerResolver := matching.NewPeerResolver(cf.resolver, namedPort)
 
-	defaultLoadBalancer := matching.NewLoadBalancer(domainIDToName, cf.dynConfig)
-	roundRobinLoadBalancer := matching.NewRoundRobinLoadBalancer(domainIDToName, cf.dynConfig)
-	weightedLoadBalancer := matching.NewWeightedLoadBalancer(roundRobinLoadBalancer, domainIDToName, cf.dynConfig, cf.logger)
+	partitionConfigProvider := matching.NewPartitionConfigProvider(cf.logger, domainIDToName, cf.dynConfig)
+	defaultLoadBalancer := matching.NewLoadBalancer(partitionConfigProvider)
+	roundRobinLoadBalancer := matching.NewRoundRobinLoadBalancer(partitionConfigProvider)
+	weightedLoadBalancer := matching.NewWeightedLoadBalancer(roundRobinLoadBalancer, partitionConfigProvider, cf.logger)
 	loadBalancers := map[string]matching.LoadBalancer{
 		"random":      defaultLoadBalancer,
 		"round-robin": roundRobinLoadBalancer,
@@ -163,6 +164,7 @@ func (cf *rpcClientFactory) NewMatchingClientWithTimeout(
 		rawClient,
 		peerResolver,
 		matching.NewMultiLoadBalancer(defaultLoadBalancer, loadBalancers, domainIDToName, cf.dynConfig, cf.logger),
+		partitionConfigProvider,
 	)
 	client = timeoutwrapper.NewMatchingClient(client, longPollTimeout, timeout)
 	if errorRate := cf.dynConfig.GetFloat64Property(dynamicconfig.MatchingErrorInjectionRate)(); errorRate != 0 {

client/matching/client.go

@@ -36,18 +36,21 @@ type clientImpl struct {
 	client       Client
 	peerResolver PeerResolver
 	loadBalancer LoadBalancer
+	provider     PartitionConfigProvider
 }
 
 // NewClient creates a new history service TChannel client
 func NewClient(
 	client Client,
 	peerResolver PeerResolver,
 	lb LoadBalancer,
+	provider PartitionConfigProvider,
 ) Client {
 	return &clientImpl{
 		client:       client,
 		peerResolver: peerResolver,
 		loadBalancer: lb,
+		provider:     provider,
 	}
 }
 
@@ -131,6 +134,12 @@ func (c *clientImpl) PollForDecisionTask(
 		return nil, err
 	}
 	request.PollRequest.TaskList.Name = originalTaskListName
+	c.provider.UpdatePartitionConfig(
+		request.GetDomainUUID(),
+		*request.PollRequest.GetTaskList(),
+		persistence.TaskListTypeDecision,
+		resp.PartitionConfig,
+	)
 	c.loadBalancer.UpdateWeight(
 		request.GetDomainUUID(),
 		*request.PollRequest.GetTaskList(),

client/matching/loadbalancer.go

@@ -28,7 +28,6 @@ import (
 	"strings"
 
 	"github.com/uber/cadence/common"
-	"github.com/uber/cadence/common/dynamicconfig"
 	"github.com/uber/cadence/common/types"
 )
 
@@ -74,22 +73,17 @@ type (
 	}
 
 	defaultLoadBalancer struct {
-		nReadPartitions  dynamicconfig.IntPropertyFnWithTaskListInfoFilters
-		nWritePartitions dynamicconfig.IntPropertyFnWithTaskListInfoFilters
-		domainIDToName   func(string) (string, error)
+		provider PartitionConfigProvider
 	}
 )
 
 // NewLoadBalancer returns an instance of matching load balancer that
 // can help distribute api calls across task list partitions
 func NewLoadBalancer(
-	domainIDToName func(string) (string, error),
-	dc *dynamicconfig.Collection,
+	provider PartitionConfigProvider,
 ) LoadBalancer {
 	return &defaultLoadBalancer{
-		domainIDToName:   domainIDToName,
-		nReadPartitions:  dc.GetIntPropertyFilteredByTaskListInfo(dynamicconfig.MatchingNumTasklistReadPartitions),
-		nWritePartitions: dc.GetIntPropertyFilteredByTaskListInfo(dynamicconfig.MatchingNumTasklistWritePartitions),
+		provider: provider,
 	}
 }
 
@@ -99,16 +93,7 @@ func (lb *defaultLoadBalancer) PickWritePartition(
 	taskListType int,
 	forwardedFrom string,
 ) string {
-	domainName, err := lb.domainIDToName(domainID)
-	if err != nil {
-		return taskList.GetName()
-	}
-	nPartitions := lb.nWritePartitions(domainName, taskList.GetName(), taskListType)
-
-	// checks to make sure number of writes never exceeds number of reads
-	if nRead := lb.nReadPartitions(domainName, taskList.GetName(), taskListType); nPartitions > nRead {
-		nPartitions = nRead
-	}
+	nPartitions := lb.provider.GetNumberOfWritePartitions(domainID, taskList, taskListType)
 	return lb.pickPartition(taskList, forwardedFrom, nPartitions)
 
 }
@@ -119,11 +104,7 @@ func (lb *defaultLoadBalancer) PickReadPartition(
 	taskListType int,
 	forwardedFrom string,
 ) string {
-	domainName, err := lb.domainIDToName(domainID)
-	if err != nil {
-		return taskList.GetName()
-	}
-	n := lb.nReadPartitions(domainName, taskList.GetName(), taskListType)
+	n := lb.provider.GetNumberOfReadPartitions(domainID, taskList, taskListType)
 	return lb.pickPartition(taskList, forwardedFrom, n)
 
 }