Set up activity throttle limit as input and propagate to server (#331)

Author: madhuravi

internal/internal_public.go

@@ -104,7 +104,8 @@ func NewActivityTaskWorker(
 		TaskList:                        taskList,
 		ConcurrentPollRoutineSize:       defaultConcurrentPollRoutineSize,
 		ConcurrentActivityExecutionSize: wOptions.MaxConcurrentActivityExecutionSize,
-		MaxActivityExecutionPerSecond:   wOptions.MaxActivityExecutionPerSecond,
+		WorkerActivitiesPerSecond:       wOptions.WorkerActivitiesPerSecond,
+		TaskListActivitiesPerSecond:     wOptions.TaskListActivitiesPerSecond,
 		Identity:                        wOptions.Identity,
 		MetricsScope:                    wOptions.MetricsScope,
 		Logger:                          wOptions.Logger,

internal/internal_task_pollers.go

@@ -81,13 +81,14 @@ type (
 
 	// activityTaskPoller implements polling/processing a workflow task
 	activityTaskPoller struct {
-		domain       string
-		taskListName string
-		identity     string
-		service      workflowserviceclient.Interface
-		taskHandler  ActivityTaskHandler
-		metricsScope tally.Scope
-		logger       *zap.Logger
+		domain              string
+		taskListName        string
+		identity            string
+		service             workflowserviceclient.Interface
+		taskHandler         ActivityTaskHandler
+		metricsScope        tally.Scope
+		logger              *zap.Logger
+		activitiesPerSecond float64
 	}
 
 	historyIteratorImpl struct {
@@ -431,13 +432,15 @@ func newGetHistoryPageFunc(
 func newActivityTaskPoller(taskHandler ActivityTaskHandler, service workflowserviceclient.Interface,
 	domain string, params workerExecutionParameters) *activityTaskPoller {
 	return &activityTaskPoller{
-		taskHandler:  taskHandler,
-		service:      metrics.NewWorkflowServiceWrapper(service, params.MetricsScope),
-		domain:       domain,
-		taskListName: params.TaskList,
-		identity:     params.Identity,
-		logger:       params.Logger,
-		metricsScope: params.MetricsScope}
+		taskHandler:         taskHandler,
+		service:             metrics.NewWorkflowServiceWrapper(service, params.MetricsScope),
+		domain:              domain,
+		taskListName:        params.TaskList,
+		identity:            params.Identity,
+		logger:              params.Logger,
+		metricsScope:        params.MetricsScope,
+		activitiesPerSecond: params.TaskListActivitiesPerSecond,
+	}
 }
 
 // Poll for a single activity task from the service
@@ -450,9 +453,10 @@ func (atp *activityTaskPoller) poll() (*activityTask, error) {
 		atp.logger.Debug("activityTaskPoller::Poll")
 	})
 	request := &s.PollForActivityTaskRequest{
-		Domain:   common.StringPtr(atp.domain),
-		TaskList: common.TaskListPtr(s.TaskList{Name: common.StringPtr(atp.taskListName)}),
-		Identity: common.StringPtr(atp.identity),
+		Domain:           common.StringPtr(atp.domain),
+		TaskList:         common.TaskListPtr(s.TaskList{Name: common.StringPtr(atp.taskListName)}),
+		Identity:         common.StringPtr(atp.identity),
+		TaskListMetadata: &s.TaskListMetadata{MaxTasksPerSecond: &atp.activitiesPerSecond},
 	}
 
 	tchCtx, cancel, opt := newChannelContext(context.Background(), chanTimeout(pollTaskServiceTimeOut))

internal/internal_worker.go

@@ -52,7 +52,9 @@ const (
 	defaultConcurrentPollRoutineSize = 2
 
 	defaultMaxConcurrentActivityExecutionSize = 1000   // Large concurrent activity execution size (1k)
-	defaultMaxActivityExecutionRate           = 100000 // Large activity execution rate (unlimited)
+	_defaultWorkerActivitiesPerSecond         = 100000 // Large activity executions/sec (unlimited)
+
+	_defaultTaskListActivitiesPerSecond = 100000.0 // Large activity executions/sec (unlimited)
 
 	defaultMaxConcurrentWorkflowExecutionSize = 50     // hardcoded max workflow execution size.
 	defaultMaxWorkflowExecutionRate           = 100000 // Large workflow execution rate (unlimited)
@@ -108,8 +110,11 @@ type (
 		// Defines how many concurrent executions for task list by this worker.
 		ConcurrentActivityExecutionSize int
 
-		// Defines rate limiting on number of activity tasks that can be executed per second.
-		MaxActivityExecutionPerSecond float64
+		// Defines rate limiting on number of activity tasks that can be executed per second per worker.
+		WorkerActivitiesPerSecond float64
+
+		// TaskListActivitiesPerSecond is the throttling limit for activity tasks controlled by the server
+		TaskListActivitiesPerSecond float64
 
 		// User can provide an identity for the debuggability. If not provided the framework has
 		// a default option.
@@ -310,7 +315,7 @@ func newActivityTaskWorker(
 		baseWorkerOptions{
 			pollerCount:       workerParams.ConcurrentPollRoutineSize,
 			maxConcurrentTask: workerParams.ConcurrentActivityExecutionSize,
-			maxTaskPerSecond:  workerParams.MaxActivityExecutionPerSecond,
+			maxTaskPerSecond:  workerParams.WorkerActivitiesPerSecond,
 			taskWorker:        poller,
 			identity:          workerParams.Identity,
 			workerType:        "ActivityWorker",
@@ -1010,14 +1015,15 @@ func newAggregatedWorker(
 		TaskList:                        taskList,
 		ConcurrentPollRoutineSize:       defaultConcurrentPollRoutineSize,
 		ConcurrentActivityExecutionSize: wOptions.MaxConcurrentActivityExecutionSize,
-		MaxActivityExecutionPerSecond:   wOptions.MaxActivityExecutionPerSecond,
+		WorkerActivitiesPerSecond:       wOptions.WorkerActivitiesPerSecond,
 		Identity:                        wOptions.Identity,
 		MetricsScope:                    wOptions.MetricsScope,
 		Logger:                          wOptions.Logger,
 		EnableLoggingInReplay:           wOptions.EnableLoggingInReplay,
 		UserContext:                     wOptions.BackgroundActivityContext,
 		DisableStickyExecution:          wOptions.DisableStickyExecution,
 		StickyScheduleToStartTimeout:    wOptions.StickyScheduleToStartTimeout,
+		TaskListActivitiesPerSecond:     wOptions.TaskListActivitiesPerSecond,
 	}
 
 	ensureRequiredParams(&workerParams)
@@ -1264,8 +1270,11 @@ func fillWorkerOptionsDefaults(options WorkerOptions) WorkerOptions {
 	if options.MaxConcurrentActivityExecutionSize == 0 {
 		options.MaxConcurrentActivityExecutionSize = defaultMaxConcurrentActivityExecutionSize
 	}
-	if options.MaxActivityExecutionPerSecond == 0 {
-		options.MaxActivityExecutionPerSecond = defaultMaxActivityExecutionRate
+	if options.WorkerActivitiesPerSecond == 0 {
+		options.WorkerActivitiesPerSecond = _defaultWorkerActivitiesPerSecond
+	}
+	if options.TaskListActivitiesPerSecond == 0 {
+		options.TaskListActivitiesPerSecond = _defaultTaskListActivitiesPerSecond
 	}
 	if options.StickyScheduleToStartTimeout.Seconds() == 0 {
 		options.StickyScheduleToStartTimeout = stickyDecisionScheduleToStartTimeoutSeconds * time.Second

internal/internal_worker_test.go

@@ -214,7 +214,7 @@ func TestCreateWorker(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	service := workflowservicetest.NewMockClient(mockCtrl)
 
-	worker := createWorker(t, service)
+	worker := createWorkerWithThrottle(t, service, float64(500.0))
 	err := worker.Start()
 	require.NoError(t, err)
 	time.Sleep(time.Millisecond * 200)
@@ -293,7 +293,33 @@ func TestWorkerStartFailsWithInvalidDomain(t *testing.T) {
 	}
 }
 
+func ofPollForActivityTaskRequest(tps float64) gomock.Matcher {
+	return &mockPollForActivityTaskRequest{tps: tps}
+}
+
+type mockPollForActivityTaskRequest struct {
+	tps float64
+}
+
+func (m *mockPollForActivityTaskRequest) Matches(x interface{}) bool {
+	v, ok := x.(*s.PollForActivityTaskRequest)
+	if !ok {
+		return false
+	}
+	return *(v.TaskListMetadata.MaxTasksPerSecond) == m.tps
+}
+
+func (m *mockPollForActivityTaskRequest) String() string {
+	return "PollForActivityTaskRequest"
+}
+
 func createWorker(t *testing.T, service *workflowservicetest.MockClient) Worker {
+	return createWorkerWithThrottle(t, service, float64(0.0))
+}
+
+func createWorkerWithThrottle(
+	t *testing.T, service *workflowservicetest.MockClient, activitiesPerSecond float64,
+) Worker {
 	domain := "testDomain"
 	domainStatus := s.DomainStatusRegistered
 	domainDesc := &s.DescribeDomainResponse{
@@ -309,7 +335,13 @@ func createWorker(t *testing.T, service *workflowservicetest.MockClient) Worker
 		}).AnyTimes()
 
 	activityTask := &s.PollForActivityTaskResponse{}
-	service.EXPECT().PollForActivityTask(gomock.Any(), gomock.Any(), callOptions...).Return(activityTask, nil).AnyTimes()
+	expectedActivitiesPerSecond := activitiesPerSecond
+	if expectedActivitiesPerSecond == 0.0 {
+		expectedActivitiesPerSecond = _defaultTaskListActivitiesPerSecond
+	}
+	service.EXPECT().PollForActivityTask(
+		gomock.Any(), ofPollForActivityTaskRequest(expectedActivitiesPerSecond), callOptions...,
+	).Return(activityTask, nil).AnyTimes()
 	service.EXPECT().RespondActivityTaskCompleted(gomock.Any(), gomock.Any(), callOptions...).Return(nil).AnyTimes()
 
 	decisionTask := &s.PollForDecisionTaskResponse{}
@@ -318,7 +350,8 @@ func createWorker(t *testing.T, service *workflowservicetest.MockClient) Worker
 
 	// Configure worker options.
 	workerOptions := WorkerOptions{}
-	workerOptions.MaxActivityExecutionPerSecond = 20
+	workerOptions.WorkerActivitiesPerSecond = 20
+	workerOptions.TaskListActivitiesPerSecond = activitiesPerSecond
 
 	// Start Worker.
 	worker := NewWorker(

internal/worker.go

@@ -49,14 +49,22 @@ type (
 		// default: defaultMaxConcurrentActivityExecutionSize(1k)
 		MaxConcurrentActivityExecutionSize int
 
-		// Optional: Sets the rate limiting on number of activities that can be executed per second. Notice that the
-		// number is represented in float, so that you can set it to less than 1 if needed. For example, set the number
-		// to 0.1 means you want your activity to be executed once for every 10 seconds. This can be used to protect
-		// down stream services from flooding.
-		// The zero value of this uses the default value.
-		// default: defaultMaxActivityExecutionRate(100k)
-		// Warning: activity's StartToCloseTimeout starts ticking even if a task is blocked due to rate limiting.
-		MaxActivityExecutionPerSecond float64
+		// Optional: Sets the rate limiting on number of activities that can be executed per second per
+		// worker. This can be used to limit resources used by the worker.
+		// Notice that the number is represented in float, so that you can set it to less than
+		// 1 if needed. For example, set the number to 0.1 means you want your activity to be executed
+		// once for every 10 seconds. This can be used to protect down stream services from flooding.
+		// The zero value of this uses the default value. Default: 100k
+		WorkerActivitiesPerSecond float64
+
+		// Optional: Sets the rate limiting on number of activities that can be executed per second.
+		// This is managed by the server and controls activities per second for your entire tasklist
+		// whereas WorkerActivityTasksPerSecond controls activities only per worker.
+		// Notice that the number is represented in float, so that you can set it to less than
+		// 1 if needed. For example, set the number to 0.1 means you want your activity to be executed
+		// once for every 10 seconds. This can be used to protect down stream services from flooding.
+		// The zero value of this uses the default value. Default: 100k
+		TaskListActivitiesPerSecond float64
 
 		// Optional: if the activities need auto heart beating for those activities
 		// by the framework