close tikv#4500: add QPS rate limiter

Signed-off-by: Cabinfever_B <cabinfeveroier@gmail.com>
CabinfeverB · Dec 27, 2021 · a1c1ecc · a1c1ecc
1 parent 8334bbe
commit a1c1ecc
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/server/self_protection.go b/server/self_protection.go
@@ -20,7 +20,8 @@ import (
 	"golang.org/x/time/rate"
 )
 
-type APIRateLimiter struct {
+// serviceRateLimiter is used to rimit service request rate
+type serviceRateLimiter struct {
 	mu sync.RWMutex
 
 	enableQPSLimit bool
@@ -33,7 +34,7 @@ type APIRateLimiter struct {
 
 // QPSAllow firstly check component token bucket and then check total token bucket
 // if component rate limiter doesn't allow, it won't ask total limiter
-func (rl *APIRateLimiter) QPSAllow(componentName string) bool {
+func (rl *serviceRateLimiter) QPSAllow(componentName string) bool {
 	if !rl.enableQPSLimit {
 		return true
 	}
@@ -53,7 +54,7 @@ func (rl *APIRateLimiter) QPSAllow(componentName string) bool {
 }
 
 // Allow currentlt only supports QPS rate limit
-func (rl *APIRateLimiter) Allow(componentName string) bool {
+func (rl *serviceRateLimiter) Allow(componentName string) bool {
 	rl.mu.RLock()
 	defer rl.mu.RUnlock()
 	return rl.QPSAllow(componentName)

diff --git a/server/self_protection_test.go b/server/self_protection_test.go
@@ -25,27 +25,27 @@ import (
 var _ = Suite(&testSelfProtectHandler{})
 
 type testSelfProtectHandler struct {
-	rateLimiterOnlyTotal         *APIRateLimiter
-	rateLimiterDisabled          *APIRateLimiter
-	rateLimiterZeroBucket        *APIRateLimiter
-	rateLimiterComopnent         *APIRateLimiter
-	rateLimiterNoComopnentConfig *APIRateLimiter
+	rateLimiterOnlyTotal         *serviceRateLimiter
+	rateLimiterDisabled          *serviceRateLimiter
+	rateLimiterZeroBucket        *serviceRateLimiter
+	rateLimiterComopnent         *serviceRateLimiter
+	rateLimiterNoComopnentConfig *serviceRateLimiter
 }
 
 func (s *testSelfProtectHandler) SetUpSuite(c *C) {
-	s.rateLimiterOnlyTotal = &APIRateLimiter{
+	s.rateLimiterOnlyTotal = &serviceRateLimiter{
 		enableQPSLimit:          true,
 		totalQPSRateLimiter:     rate.NewLimiter(100, 100),
 		enableComponentQPSLimit: false,
 	}
-	s.rateLimiterDisabled = &APIRateLimiter{
+	s.rateLimiterDisabled = &serviceRateLimiter{
 		enableQPSLimit: false,
 	}
-	s.rateLimiterZeroBucket = &APIRateLimiter{
+	s.rateLimiterZeroBucket = &serviceRateLimiter{
 		enableQPSLimit:      true,
 		totalQPSRateLimiter: rate.NewLimiter(0, 0),
 	}
-	s.rateLimiterComopnent = &APIRateLimiter{
+	s.rateLimiterComopnent = &serviceRateLimiter{
 		enableQPSLimit:          true,
 		totalQPSRateLimiter:     rate.NewLimiter(100, 100),
 		enableComponentQPSLimit: true,
@@ -54,7 +54,7 @@ func (s *testSelfProtectHandler) SetUpSuite(c *C) {
 	s.rateLimiterComopnent.componentQPSRateLimiter["pdctl"] = rate.NewLimiter(100, 100)
 	s.rateLimiterComopnent.componentQPSRateLimiter["anonymous"] = rate.NewLimiter(100, 100)
 
-	s.rateLimiterNoComopnentConfig = &APIRateLimiter{
+	s.rateLimiterNoComopnentConfig = &serviceRateLimiter{
 		enableQPSLimit:          true,
 		totalQPSRateLimiter:     rate.NewLimiter(200, 200),
 		enableComponentQPSLimit: true,
@@ -63,7 +63,7 @@ func (s *testSelfProtectHandler) SetUpSuite(c *C) {
 	s.rateLimiterNoComopnentConfig.componentQPSRateLimiter["pdctl"] = rate.NewLimiter(10, 10)
 }
 
-func CountRateLimiterHandleResult(handler *APIRateLimiter, component string, successCount *int,
+func CountRateLimiterHandleResult(handler *serviceRateLimiter, component string, successCount *int,
 	failedCount *int, lock *sync.Mutex, wg *sync.WaitGroup) {
 	result := handler.Allow(component)
 	lock.Lock()