fix rate limiter bug (#230)

* fix rate limiter bug

* add comments on how to set rate to less than 1

Author: yiminc

internal_public.go

@@ -94,16 +94,15 @@ func NewActivityTaskWorker(
 ) Worker {
 	wOptions := fillWorkerOptionsDefaults(options)
 	workerParams := workerExecutionParameters{
-		TaskList:                                taskList,
-		ConcurrentPollRoutineSize:               defaultConcurrentPollRoutineSize,
-		ConcurrentActivityExecutionSize:         wOptions.MaxConcurrentActivityExecutionSize,
-		MaxActivityExecutionRate:                wOptions.MaxActivityExecutionRate,
-		MaxActivityExecutionRateRefreshDuration: wOptions.MaxActivityExecutionRateRefreshDuration,
-		Identity:              wOptions.Identity,
-		MetricsScope:          wOptions.MetricsScope,
-		Logger:                wOptions.Logger,
-		EnableLoggingInReplay: wOptions.EnableLoggingInReplay,
-		UserContext:           wOptions.BackgroundActivityContext,
+		TaskList:                        taskList,
+		ConcurrentPollRoutineSize:       defaultConcurrentPollRoutineSize,
+		ConcurrentActivityExecutionSize: wOptions.MaxConcurrentActivityExecutionSize,
+		MaxActivityExecutionPerSecond:   wOptions.MaxActivityExecutionPerSecond,
+		Identity:                        wOptions.Identity,
+		MetricsScope:                    wOptions.MetricsScope,
+		Logger:                          wOptions.Logger,
+		EnableLoggingInReplay:           wOptions.EnableLoggingInReplay,
+		UserContext:                     wOptions.BackgroundActivityContext,
 	}
 
 	processTestTags(&wOptions, &workerParams)

internal_worker.go

@@ -33,7 +33,6 @@ import (
 	"strconv"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/apache/thrift/lib/go/thrift"
 	"github.com/uber-go/tally"
@@ -107,11 +106,8 @@ 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 MaxActivityExecutionRateRefreshDuration.
-		MaxActivityExecutionRate int
-
-		// Defines refresh duration for rate limit. If not specified, it uses 1s as default.
-		MaxActivityExecutionRateRefreshDuration time.Duration
+		// Defines rate limiting on number of activity tasks that can be executed per second.
+		MaxActivityExecutionPerSecond float64
 
 		// User can provide an identity for the debuggability. If not provided the framework has
 		// a default option.
@@ -226,13 +222,12 @@ func newWorkflowTaskWorkerInternal(
 		params,
 	)
 	worker := newBaseWorker(baseWorkerOptions{
-		pollerCount:                params.ConcurrentPollRoutineSize,
-		maxConcurrentTask:          defaultMaxConcurrentWorkflowExecutionSize,
-		maxTaskRate:                defaultMaxWorkflowExecutionRate,
-		maxTaskRateRefreshDuration: time.Second,
-		taskWorker:                 poller,
-		identity:                   params.Identity,
-		workerType:                 "DecisionWorker"},
+		pollerCount:       params.ConcurrentPollRoutineSize,
+		maxConcurrentTask: defaultMaxConcurrentWorkflowExecutionSize,
+		maxTaskPerSecond:  defaultMaxWorkflowExecutionRate,
+		taskWorker:        poller,
+		identity:          params.Identity,
+		workerType:        "DecisionWorker"},
 		params.Logger,
 		params.MetricsScope,
 	)
@@ -306,13 +301,12 @@ func newActivityTaskWorker(
 
 	base := newBaseWorker(
 		baseWorkerOptions{
-			pollerCount:                workerParams.ConcurrentPollRoutineSize,
-			maxConcurrentTask:          workerParams.ConcurrentActivityExecutionSize,
-			maxTaskRate:                workerParams.MaxActivityExecutionRate,
-			maxTaskRateRefreshDuration: workerParams.MaxActivityExecutionRateRefreshDuration,
-			taskWorker:                 poller,
-			identity:                   workerParams.Identity,
-			workerType:                 "ActivityWorker",
+			pollerCount:       workerParams.ConcurrentPollRoutineSize,
+			maxConcurrentTask: workerParams.ConcurrentActivityExecutionSize,
+			maxTaskPerSecond:  workerParams.MaxActivityExecutionPerSecond,
+			taskWorker:        poller,
+			identity:          workerParams.Identity,
+			workerType:        "ActivityWorker",
 		},
 		workerParams.Logger,
 		workerParams.MetricsScope,
@@ -1005,16 +999,15 @@ func newAggregatedWorker(
 ) (worker Worker) {
 	wOptions := fillWorkerOptionsDefaults(options)
 	workerParams := workerExecutionParameters{
-		TaskList:                                taskList,
-		ConcurrentPollRoutineSize:               defaultConcurrentPollRoutineSize,
-		ConcurrentActivityExecutionSize:         wOptions.MaxConcurrentActivityExecutionSize,
-		MaxActivityExecutionRate:                wOptions.MaxActivityExecutionRate,
-		MaxActivityExecutionRateRefreshDuration: wOptions.MaxActivityExecutionRateRefreshDuration,
-		Identity:              wOptions.Identity,
-		MetricsScope:          wOptions.MetricsScope,
-		Logger:                wOptions.Logger,
-		EnableLoggingInReplay: wOptions.EnableLoggingInReplay,
-		UserContext:           wOptions.BackgroundActivityContext,
+		TaskList:                        taskList,
+		ConcurrentPollRoutineSize:       defaultConcurrentPollRoutineSize,
+		ConcurrentActivityExecutionSize: wOptions.MaxConcurrentActivityExecutionSize,
+		MaxActivityExecutionPerSecond:   wOptions.MaxActivityExecutionPerSecond,
+		Identity:                        wOptions.Identity,
+		MetricsScope:                    wOptions.MetricsScope,
+		Logger:                          wOptions.Logger,
+		EnableLoggingInReplay:           wOptions.EnableLoggingInReplay,
+		UserContext:                     wOptions.BackgroundActivityContext,
 	}
 
 	ensureRequiredParams(&workerParams)
@@ -1227,11 +1220,8 @@ func fillWorkerOptionsDefaults(options WorkerOptions) WorkerOptions {
 	if options.MaxConcurrentActivityExecutionSize == 0 {
 		options.MaxConcurrentActivityExecutionSize = defaultMaxConcurrentActivityExecutionSize
 	}
-	if options.MaxActivityExecutionRate == 0 {
-		options.MaxActivityExecutionRate = defaultMaxActivityExecutionRate
-	}
-	if options.MaxActivityExecutionRateRefreshDuration == 0 {
-		options.MaxActivityExecutionRateRefreshDuration = time.Second
+	if options.MaxActivityExecutionPerSecond == 0 {
+		options.MaxActivityExecutionPerSecond = defaultMaxActivityExecutionRate
 	}
 	return options
 }

internal_worker_base.go

@@ -82,13 +82,12 @@ type (
 
 	// baseWorkerOptions options to configure base worker.
 	baseWorkerOptions struct {
-		pollerCount                int
-		maxConcurrentTask          int
-		maxTaskRate                int
-		maxTaskRateRefreshDuration time.Duration
-		taskWorker                 taskPoller
-		identity                   string
-		workerType                 string
+		pollerCount       int
+		maxConcurrentTask int
+		maxTaskPerSecond  float64
+		taskWorker        taskPoller
+		identity          string
+		workerType        string
 	}
 
 	// baseWorker that wraps worker activities.
@@ -127,8 +126,8 @@ func newBaseWorker(options baseWorkerOptions, logger *zap.Logger, metricsScope t
 	return &baseWorker{
 		options:         options,
 		shutdownCh:      make(chan struct{}),
-		pollLimiter:     rate.NewLimiter(rate.Every(time.Millisecond*100), 100),
-		taskLimiter:     rate.NewLimiter(rate.Every(options.maxTaskRateRefreshDuration), options.maxTaskRate),
+		pollLimiter:     rate.NewLimiter(rate.Limit(1000), 1),
+		taskLimiter:     rate.NewLimiter(rate.Limit(options.maxTaskPerSecond), 1),
 		retrier:         backoff.NewConcurrentRetrier(pollOperationRetryPolicy),
 		logger:          logger.With(zapcore.Field{Key: tagWorkerType, Type: zapcore.StringType, String: options.workerType}),
 		metricsScope:    tagScope(metricsScope, tagWorkerType, options.workerType),
@@ -160,8 +159,7 @@ func (bw *baseWorker) Start() {
 		bw.logger.Info("Started Worker",
 			zap.Int("PollerCount", bw.options.pollerCount),
 			zap.Int("MaxConcurrentTask", bw.options.maxConcurrentTask),
-			zap.Int("MaxTaskRate", bw.options.maxTaskRate),
-			zap.Duration("MaxTaskRateRefreshDuration", bw.options.maxTaskRateRefreshDuration),
+			zap.Float64("MaxTaskPerSecond", bw.options.maxTaskPerSecond),
 		)
 	})
 }

internal_worker_test.go

@@ -340,7 +340,7 @@ func createWorker(t *testing.T, service *mocks.TChanWorkflowService) Worker {
 
 	// Configure worker options.
 	workerOptions := WorkerOptions{}
-	workerOptions.MaxActivityExecutionRate = 20
+	workerOptions.MaxActivityExecutionPerSecond = 20
 
 	// Start Worker.
 	worker := NewWorker(

worker.go

@@ -22,7 +22,6 @@ package cadence
 
 import (
 	"context"
-	"time"
 
 	"github.com/uber-go/tally"
 
@@ -53,18 +52,14 @@ type (
 		// default: defaultMaxConcurrentActivityExecutionSize(1k)
 		MaxConcurrentActivityExecutionSize int
 
-		// Optional: Sets the rate limiting on number of activities that can be executed per refresh duration.
-		// This can be used to protect down stream services from flooding.
+		// 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.
-		MaxActivityExecutionRate int
-
-		// Optional: Sets the refresh duration for rate limit. If not specified, it uses 1s as default.
-		// Use this to fine tune your rate limiter. For example, you could set this duration to 10s with max rate set to
-		// 1 which means we rate limit the activity task to at most 1 per every 10s. You could also set this duration to
-		// 100ms with max rate set wo 10 which means rate limit to 10 activity task execution per every 100ms.
-		MaxActivityExecutionRateRefreshDuration time.Duration
+		MaxActivityExecutionPerSecond float64
 
 		// Optional: if the activities need auto heart beating for those activities
 		// by the framework