Use dynamic rate limiter for workflow handler (#2213)

Update the rate limiter to take in an RPSFunc to handle changing RPS information

common/quotas/interfaces.go

@@ -20,6 +20,9 @@
 
 package quotas
 
+// RPSFunc returns a float64 as the RPS
+type RPSFunc func() float64
+
 // Policy corresponds to a quota policy. A policy allows implementing layered
 // and more complex rate limiting functionality.
 type Policy interface {

common/quotas/ratelimiter.go

@@ -30,6 +30,8 @@ import (
 	"golang.org/x/time/rate"
 )
 
+const _defaultRPSTTL = 60 * time.Second
+
 // RateLimiter is a wrapper around the golang rate limiter handling dynamic
 // configuration updates of the max dispatch per second. This has comparable
 // performance to the token bucket rate limiter.
@@ -123,3 +125,21 @@ func (rl *RateLimiter) shouldUpdate(maxDispatchPerSecond *float64) bool {
 		return *maxDispatchPerSecond < *rl.maxDispatchPerSecond
 	}
 }
+
+type dynamicRateLimiter struct {
+	rps RPSFunc
+	rl  *RateLimiter
+}
+
+// NewDynamicRateLimiter returns a rate limiter which handles dynamic config
+func NewDynamicRateLimiter(rps RPSFunc) Policy {
+	initialRps := rps()
+	rl := NewRateLimiter(&initialRps, _defaultRPSTTL, int(rps()))
+	return &dynamicRateLimiter{rps, rl}
+}
+
+func (d *dynamicRateLimiter) Allow() bool {
+	rps := float64(d.rps())
+	d.rl.UpdateMaxDispatch(&rps)
+	return d.rl.Allow()
+}

common/tokenbucket/tb.go

@@ -140,19 +140,6 @@ func newTokenBucket(rps int, timeSource clock.TimeSource) *tokenBucketImpl {
 	return tb
 }
 
-// NewFactory creates an instance of factory used for creating TokenBucket instances
-func NewFactory() Factory {
-	return &tokenBucketFactoryImpl{}
-}
-
-// CreateTokenBucket creates and returns a
-// new token bucket rate limiter that
-// repelenishes the bucket every 100
-// milliseconds. Thread safe.
-func (f *tokenBucketFactoryImpl) CreateTokenBucket(rps int, timeSource clock.TimeSource) TokenBucket {
-	return New(rps, timeSource)
-}
-
 func (tb *tokenBucketImpl) TryConsume(count int) (bool, time.Duration) {
 	now := tb.timeSource.Now().UnixNano()
 	tb.Lock()

service/frontend/workflowHandler.go

@@ -44,7 +44,6 @@ import (
 	"github.com/uber/cadence/common/blobstore"
 	"github.com/uber/cadence/common/cache"
 	"github.com/uber/cadence/common/client"
-	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/elasticsearch/validator"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/tag"
@@ -53,7 +52,6 @@ import (
 	"github.com/uber/cadence/common/persistence"
 	"github.com/uber/cadence/common/quotas"
 	"github.com/uber/cadence/common/service"
-	"github.com/uber/cadence/common/tokenbucket"
 	warchiver "github.com/uber/cadence/service/worker/archiver"
 	"go.uber.org/yarpc/yarpcerrors"
 )
@@ -165,7 +163,9 @@ func NewWorkflowHandler(sVice service.Service, config *Config, metadataMgr persi
 		tokenSerializer: common.NewJSONTaskTokenSerializer(),
 		metricsClient:   sVice.GetMetricsClient(),
 		domainCache:     cache.NewDomainCache(metadataMgr, sVice.GetClusterMetadata(), sVice.GetMetricsClient(), sVice.GetLogger()),
-		rateLimiter:     quotas.NewSimpleRateLimiter(tokenbucket.NewDynamicTokenBucket(config.RPS, clock.NewRealTimeSource())),
+		rateLimiter: quotas.NewDynamicRateLimiter(func() float64 {
+			return float64(config.RPS())
+		}),
 		blobstoreClient: blobstoreClient,
 		versionChecker:  &versionChecker{checkVersion: config.EnableClientVersionCheck()},
 		domainHandler: newDomainHandler(