Increase number of forward tokens for isolation tasklists

uber · Jun 2, 2023 · a547562 · a547562
1 parent d17aff4
commit a547562
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 12 deletions.
diff --git a/service/matching/forwarder.go b/service/matching/forwarder.go
@@ -57,6 +57,8 @@ type (
 		// todo: implement a rate limiter that automatically
 		// adjusts rate based on ServiceBusy errors from API calls
 		limiter *quotas.DynamicRateLimiter
+
+		isolationGroups []string
 	}
 	// ForwarderReqToken is the token that must be acquired before
 	// making forwarder API calls. This type contains the state
@@ -91,19 +93,21 @@ func newForwarder(
 	taskListID *taskListID,
 	kind types.TaskListKind,
 	client matching.Client,
+	isolationGroups []string,
 ) *Forwarder {
 	rpsFunc := func() float64 { return float64(cfg.ForwarderMaxRatePerSecond()) }
 	fwdr := &Forwarder{
 		cfg:                   cfg,
 		client:                client,
 		taskListID:            taskListID,
 		taskListKind:          kind,
-		outstandingTasksLimit: int32(cfg.ForwarderMaxOutstandingTasks()),
-		outstandingPollsLimit: int32(cfg.ForwarderMaxOutstandingPolls()),
+		outstandingTasksLimit: int32(cfg.ForwarderMaxOutstandingTasks() * (len(isolationGroups) + 1)),
+		outstandingPollsLimit: int32(cfg.ForwarderMaxOutstandingPolls() * (len(isolationGroups) + 1)),
 		limiter:               quotas.NewDynamicRateLimiter(rpsFunc),
+		isolationGroups:       isolationGroups,
 	}
-	fwdr.addReqToken.Store(newForwarderReqToken(cfg.ForwarderMaxOutstandingTasks()))
-	fwdr.pollReqToken.Store(newForwarderReqToken(cfg.ForwarderMaxOutstandingPolls()))
+	fwdr.addReqToken.Store(newForwarderReqToken(int(fwdr.outstandingTasksLimit)))
+	fwdr.pollReqToken.Store(newForwarderReqToken(int(fwdr.outstandingPollsLimit)))
 	return fwdr
 }
 
@@ -250,15 +254,16 @@ func (fwdr *Forwarder) ForwardPoll(ctx context.Context) (*InternalTask, error) {
 // that's necessary before making a ForwardTask or ForwardQueryTask API call.
 // After the API call is invoked, token.release() must be invoked
 func (fwdr *Forwarder) AddReqTokenC() <-chan *ForwarderReqToken {
-	fwdr.refreshTokenC(&fwdr.addReqToken, &fwdr.outstandingTasksLimit, int32(fwdr.cfg.ForwarderMaxOutstandingTasks()))
+	fwdr.refreshTokenC(&fwdr.addReqToken, &fwdr.outstandingTasksLimit, int32(fwdr.cfg.ForwarderMaxOutstandingTasks()*(len(fwdr.isolationGroups)+1)))
 	return fwdr.addReqToken.Load().(*ForwarderReqToken).ch
 }
 
 // PollReqTokenC returns a channel that can be used to wait for a token
 // that's necessary before making a ForwardPoll API call. After the API
 // call is invoked, token.release() must be invoked
+// TODO: consider having separate token pools for different isolation groups
 func (fwdr *Forwarder) PollReqTokenC() <-chan *ForwarderReqToken {
-	fwdr.refreshTokenC(&fwdr.pollReqToken, &fwdr.outstandingPollsLimit, int32(fwdr.cfg.ForwarderMaxOutstandingPolls()))
+	fwdr.refreshTokenC(&fwdr.pollReqToken, &fwdr.outstandingPollsLimit, int32(fwdr.cfg.ForwarderMaxOutstandingPolls()*(len(fwdr.isolationGroups)+1)))
 	return fwdr.pollReqToken.Load().(*ForwarderReqToken).ch
 }
 

diff --git a/service/matching/forwarder_test.go b/service/matching/forwarder_test.go
@@ -63,7 +63,7 @@ func (t *ForwarderTestSuite) SetupTest() {
 		ForwarderMaxOutstandingTasks: func() int { return 1 },
 	}
 	t.taskList = newTestTaskListID("fwdr", "tl0", persistence.TaskListTypeDecision)
-	t.fwdr = newForwarder(t.cfg, t.taskList, types.TaskListKindNormal, t.client)
+	t.fwdr = newForwarder(t.cfg, t.taskList, types.TaskListKindNormal, t.client, nil)
 }
 
 func (t *ForwarderTestSuite) TearDownTest() {

diff --git a/service/matching/matcher_test.go b/service/matching/matcher_test.go
@@ -69,7 +69,7 @@ func (t *MatcherTestSuite) SetupTest() {
 		ForwarderMaxChildrenPerNode:  func() int { return 20 },
 	}
 	t.cfg = tlCfg
-	t.fwdr = newForwarder(&t.cfg.forwarderConfig, t.taskList, types.TaskListKindNormal, t.client)
+	t.fwdr = newForwarder(&t.cfg.forwarderConfig, t.taskList, types.TaskListKindNormal, t.client, nil)
 	t.matcher = newTaskMatcher(tlCfg, t.fwdr, metrics.NoopScope(metrics.Matching), []string{"dca1", "dca2"})
 
 	rootTaskList := newTestTaskListID(t.taskList.domainID, t.taskList.Parent(20), persistence.TaskListTypeDecision)

diff --git a/service/matching/taskListManager.go b/service/matching/taskListManager.go
@@ -189,14 +189,14 @@ func newTaskListManager(
 	tlMgr.liveness = newLiveness(clock.NewRealTimeSource(), taskListConfig.IdleTasklistCheckInterval(), tlMgr.Stop)
 	tlMgr.taskWriter = newTaskWriter(tlMgr)
 	tlMgr.taskReader = newTaskReader(tlMgr)
-	var fwdr *Forwarder
-	if tlMgr.isFowardingAllowed(taskList, *taskListKind) {
-		fwdr = newForwarder(&taskListConfig.forwarderConfig, taskList, *taskListKind, e.matchingClient)
-	}
 	var isolationGroups []string
 	if tlMgr.isIsolationMatcherEnabled() {
 		isolationGroups = config.AllIsolationGroups
 	}
+	var fwdr *Forwarder
+	if tlMgr.isFowardingAllowed(taskList, *taskListKind) {
+		fwdr = newForwarder(&taskListConfig.forwarderConfig, taskList, *taskListKind, e.matchingClient, isolationGroups)
+	}
 	tlMgr.matcher = newTaskMatcher(taskListConfig, fwdr, tlMgr.scope, isolationGroups)
 	tlMgr.startWG.Add(1)
 	return tlMgr, nil