[servicecomb-java-chassis] 02/04: [SCB-2043] flow control support custom

[li...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

liubao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/servicecomb-java-chassis.git

commit e414f5709efa4f6b0e464a62aab8235d26c94d8f
Author: GuoYL <gy...@gmail.com>
AuthorDate: Mon Jul 27 19:34:44 2020 +0800

    [SCB-2043] flow control support custom
---
 .../main/java/org/apache/servicecomb/qps/QpsControllerManager.java  | 6 ++++++
 .../main/java/org/apache/servicecomb/qps/strategy/StrategyType.java | 3 ++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/QpsControllerManager.java b/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/QpsControllerManager.java
index 40e25c5..606b642 100644
--- a/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/QpsControllerManager.java
+++ b/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/QpsControllerManager.java
@@ -22,6 +22,7 @@ import java.util.Map.Entry;
 
 import org.apache.servicecomb.core.Invocation;
 import org.apache.servicecomb.foundation.common.concurrent.ConcurrentHashMapEx;
+import org.apache.servicecomb.foundation.common.utils.SPIServiceUtils;
 import org.apache.servicecomb.qps.strategy.AbstractQpsStrategy;
 import org.apache.servicecomb.qps.strategy.FixedWindowStrategy;
 import org.apache.servicecomb.qps.strategy.LeakyBucketStrategy;
@@ -211,6 +212,8 @@ public class QpsControllerManager {
   private AbstractQpsStrategy chooseStrategy(String globalConfigKey, Long limit, Long bucket,
       String strategyKey) {
     AbstractQpsStrategy strategy;
+    AbstractQpsStrategy customStrategy = SPIServiceUtils
+        .getTargetService(AbstractQpsStrategy.class);
     switch (StrategyType.parseStrategyType(strategyKey)) {
       case FixedWindow:
         strategy = new FixedWindowStrategy(globalConfigKey, limit);
@@ -221,6 +224,9 @@ public class QpsControllerManager {
       case TokenBucket:
         strategy = new TokenBucketStrategy(globalConfigKey, limit, bucket);
         break;
+      case Custom:
+        strategy = customStrategy;
+        break;
       case SlidingWindow:
       default:
         strategy = new FixedWindowStrategy(globalConfigKey, limit);
diff --git a/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/strategy/StrategyType.java b/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/strategy/StrategyType.java
index 98d3764..7412b0a 100644
--- a/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/strategy/StrategyType.java
+++ b/handlers/handler-flowcontrol-qps/src/main/java/org/apache/servicecomb/qps/strategy/StrategyType.java
@@ -23,7 +23,8 @@ public enum StrategyType {
   TokenBucket,
   LeakyBucket,
   FixedWindow,
-  SlidingWindow;
+  SlidingWindow,
+  Custom;
 
 
   public static StrategyType parseStrategyType(String type) {