From a1c1eccefcd1ec8847ee5d6ee4027739aa4d47a5 Mon Sep 17 00:00:00 2001
From: Cabinfever_B <cabinfeveroier@gmail.com>
Date: Mon, 27 Dec 2021 11:43:29 +0800
Subject: [PATCH] close #4500: add QPS rate limiter

Signed-off-by: Cabinfever_B <cabinfeveroier@gmail.com>
---
 server/self_protection.go      |  7 ++++---
 server/self_protection_test.go | 22 +++++++++++-----------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/server/self_protection.go b/server/self_protection.go
index 89b115a4d77..d2f40d60012 100644
--- 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
index f5d6a658a42..65a909b9625 100644
--- 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()