Add per domain limits dynamic config (cadence-workflow#2277)

MirrorFork · Jul 25, 2019 · f37857a · f37857a
1 parent 3197e1b
commit f37857a
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 8 deletions.
diff --git a/common/quotas/interfaces.go b/common/quotas/interfaces.go
@@ -25,6 +25,9 @@ import "context"
 // RPSFunc returns a float64 as the RPS
 type RPSFunc func() float64
 
+// RPSKeyFunc returns a float64 as the RPS for the given key
+type RPSKeyFunc func(key string) float64
+
 // Info corresponds to information required to determine rate limits
 type Info struct {
 	Domain string

diff --git a/common/quotas/limiter_test.go b/common/quotas/limiter_test.go
@@ -89,7 +89,7 @@ func getPolicy() Policy {
 		func() float64 {
 			return float64(defaultRps)
 		},
-		func() float64 {
+		func(domain string) float64 {
 			return float64(defaultRps)
 		},
 	)

diff --git a/common/quotas/policy.go b/common/quotas/policy.go
@@ -28,14 +28,14 @@ import (
 type MultiStageRateLimiter struct {
 	sync.RWMutex
 	rps            RPSFunc
-	domainRPS      RPSFunc
+	domainRPS      RPSKeyFunc
 	domainLimiters map[string]*DynamicRateLimiter
 	globalLimiter  *DynamicRateLimiter
 }
 
 // NewMultiStageRateLimiter returns a new domain quota rate limiter. This is about
 // an order of magnitude slower than
-func NewMultiStageRateLimiter(rps RPSFunc, domainRps RPSFunc) *MultiStageRateLimiter {
+func NewMultiStageRateLimiter(rps RPSFunc, domainRps RPSKeyFunc) *MultiStageRateLimiter {
 	rl := &MultiStageRateLimiter{
 		rps:            rps,
 		domainRPS:      domainRps,
@@ -62,7 +62,11 @@ func (d *MultiStageRateLimiter) Allow(info Info) bool {
 
 	if !ok {
 		// create a new limiter
-		domainLimiter := NewDynamicRateLimiter(d.domainRPS)
+		domainLimiter := NewDynamicRateLimiter(
+			func() float64 {
+				return d.domainRPS(domain)
+			},
+		)
 
 		// verify that it is needed and add to map
 		d.Lock()

diff --git a/service/frontend/service.go b/service/frontend/service.go
@@ -51,7 +51,7 @@ type Config struct {
 	ESIndexMaxResultWindow          dynamicconfig.IntPropertyFn
 	HistoryMaxPageSize              dynamicconfig.IntPropertyFnWithDomainFilter
 	RPS                             dynamicconfig.IntPropertyFn
-	DomainRPS                       dynamicconfig.IntPropertyFn
+	DomainRPS                       dynamicconfig.IntPropertyFnWithDomainFilter
 	MaxIDLengthLimit                dynamicconfig.IntPropertyFn
 	EnableClientVersionCheck        dynamicconfig.BoolPropertyFn
 	MinRetentionDays                dynamicconfig.IntPropertyFn
@@ -97,7 +97,7 @@ func NewConfig(dc *dynamicconfig.Collection, numHistoryShards int, enableVisibil
 		ESIndexMaxResultWindow:              dc.GetIntProperty(dynamicconfig.FrontendESIndexMaxResultWindow, 10000),
 		HistoryMaxPageSize:                  dc.GetIntPropertyFilteredByDomain(dynamicconfig.FrontendHistoryMaxPageSize, common.GetHistoryMaxPageSize),
 		RPS:                                 dc.GetIntProperty(dynamicconfig.FrontendRPS, 1200),
-		DomainRPS:                           dc.GetIntProperty(dynamicconfig.FrontendDomainRPS, 1200),
+		DomainRPS:                           dc.GetIntPropertyFilteredByDomain(dynamicconfig.FrontendDomainRPS, 1200),
 		MaxIDLengthLimit:                    dc.GetIntProperty(dynamicconfig.MaxIDLengthLimit, 1000),
 		HistoryMgrNumConns:                  dc.GetIntProperty(dynamicconfig.FrontendHistoryMgrNumConns, 10),
 		MaxBadBinaries:                      dc.GetIntPropertyFilteredByDomain(dynamicconfig.FrontendMaxBadBinaries, 10),

diff --git a/service/frontend/workflowHandler.go b/service/frontend/workflowHandler.go
@@ -170,8 +170,8 @@ func NewWorkflowHandler(
 			func() float64 {
 				return float64(config.RPS())
 			},
-			func() float64 {
-				return float64(config.DomainRPS())
+			func(domain string) float64 {
+				return float64(config.DomainRPS(domain))
 			},
 		),
 		versionChecker: &versionChecker{checkVersion: config.EnableClientVersionCheck()},