support custom traffic shaping controller

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleUtil.java

@@ -20,15 +20,18 @@
 import com.alibaba.csp.sentinel.slots.block.RuleConstant;
 import com.alibaba.csp.sentinel.slots.block.RuleManager;
 import com.alibaba.csp.sentinel.slots.block.flow.controller.DefaultController;
-import com.alibaba.csp.sentinel.slots.block.flow.controller.ThrottlingController;
-import com.alibaba.csp.sentinel.slots.block.flow.controller.WarmUpController;
-import com.alibaba.csp.sentinel.slots.block.flow.controller.WarmUpRateLimiterController;
 import com.alibaba.csp.sentinel.util.StringUtil;
 import com.alibaba.csp.sentinel.util.function.Function;
 import com.alibaba.csp.sentinel.util.function.Predicate;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -132,18 +135,9 @@ public static <K> Map<K, List<FlowRule>> buildFlowRuleMap(List<FlowRule> list, F
 
     private static TrafficShapingController generateRater(/*@Valid*/ FlowRule rule) {
         if (rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {
-            switch (rule.getControlBehavior()) {
-                case RuleConstant.CONTROL_BEHAVIOR_WARM_UP:
-                    return new WarmUpController(rule.getCount(), rule.getWarmUpPeriodSec(),
-                            ColdFactorProperty.coldFactor);
-                case RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER:
-                    return new ThrottlingController(rule.getMaxQueueingTimeMs(), rule.getCount());
-                case RuleConstant.CONTROL_BEHAVIOR_WARM_UP_RATE_LIMITER:
-                    return new WarmUpRateLimiterController(rule.getCount(), rule.getWarmUpPeriodSec(),
-                            rule.getMaxQueueingTimeMs(), ColdFactorProperty.coldFactor);
-                case RuleConstant.CONTROL_BEHAVIOR_DEFAULT:
-                default:
-                    // Default mode or unknown mode: default traffic shaping controller (fast-reject).
+            TrafficShapingControllerFactory trafficShapingControllerFactory = TrafficShapingControllerFactories.get(rule.getControlBehavior());
+            if (trafficShapingControllerFactory != null) {
+                return trafficShapingControllerFactory.create(rule);
             }
         }
         return new DefaultController(rule.getCount(), rule.getGrade());

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/LoggingTrafficShapingControllerFactory.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow;
+
+import com.alibaba.csp.sentinel.log.RecordLog;
+
+import java.util.Objects;
+
+public class LoggingTrafficShapingControllerFactory implements TrafficShapingControllerFactory {
+
+    private final TrafficShapingControllerFactory delegate;
+
+    public LoggingTrafficShapingControllerFactory(TrafficShapingControllerFactory delegate) {
+        this.delegate = Objects.requireNonNull(delegate, "delegate must be not null.");
+    }
+
+    @Override
+    public TrafficShapingController create(FlowRule rule) {
+        RecordLog.debug("Creating traffic shaping controller '" + delegate.getClass().getName() + "' for rule: " + rule);
+        return delegate.create(rule);
+    }
+
+    @Override
+    public int getControlBehavior() {
+        return delegate.getControlBehavior();
+    }
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactories.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow;
+
+import com.alibaba.csp.sentinel.spi.SpiLoader;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.BinaryOperator;
+import java.util.stream.Collectors;
+
+public class TrafficShapingControllerFactories {
+
+    private static final Map<Integer, TrafficShapingControllerFactory> FACTORIES;
+
+    static {
+        FACTORIES = initFactories();
+    }
+
+    /**
+     * Using existing factory if the factory with the same control behavior already exists, because the existing factory has higher priority
+     */
+    private static BinaryOperator<TrafficShapingControllerFactory> usingExisting() {
+        return (existing, replacement) -> existing;
+    }
+
+    private static TrafficShapingControllerFactory logging(TrafficShapingControllerFactory factory) {
+        return new LoggingTrafficShapingControllerFactory(factory);
+    }
+
+    private static Map<Integer, TrafficShapingControllerFactory> initFactories() {
+        return SpiLoader.of(TrafficShapingControllerFactory.class)
+                .loadInstanceListSorted()
+                .stream()
+                .collect(Collectors.toMap(TrafficShapingControllerFactory::getControlBehavior,
+                        TrafficShapingControllerFactories::logging,
+                        usingExisting(),
+                        HashMap::new));
+    }
+
+    public static TrafficShapingControllerFactory get(int controlBehavior) {
+        return FACTORIES.get(controlBehavior);
+    }
+
+    private TrafficShapingControllerFactories() {
+    }
+
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactory.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow;
+
+/**
+ * a factory interface to create TrafficShapingController instance
+ */
+public interface TrafficShapingControllerFactory {
+
+    /**
+     * create a TrafficShapingController instance
+     * @param rule flow rule
+     * @return a new TrafficShapingController instance
+     */
+    TrafficShapingController create(FlowRule rule);
+
+    /**
+     * get the factory control behavior
+     * @return the control behavior
+     */
+    int getControlBehavior();
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/DefaultControllerFactory.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.controller;
+
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingController;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory;
+
+public class DefaultControllerFactory implements TrafficShapingControllerFactory {
+
+    @Override
+    public TrafficShapingController create(FlowRule rule) {
+        return new DefaultController(rule.getCount(), rule.getGrade());
+    }
+
+    @Override
+    public int getControlBehavior() {
+        return RuleConstant.CONTROL_BEHAVIOR_DEFAULT;
+    }
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/ThrottlingControllerFactory.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.controller;
+
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingController;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory;
+
+public class ThrottlingControllerFactory implements TrafficShapingControllerFactory {
+
+    @Override
+    public TrafficShapingController create(FlowRule rule) {
+        return new ThrottlingController(rule.getMaxQueueingTimeMs(), rule.getCount());
+    }
+
+    @Override
+    public int getControlBehavior() {
+        return RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER;
+    }
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpControllerFactory.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.controller;
+
+import com.alibaba.csp.sentinel.config.SentinelConfig;
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingController;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory;
+
+public class WarmUpControllerFactory implements TrafficShapingControllerFactory {
+
+    static volatile int coldFactor = SentinelConfig.coldFactor();
+
+    @Override
+    public TrafficShapingController create(FlowRule rule) {
+        return new WarmUpController(rule.getCount(), rule.getWarmUpPeriodSec(), coldFactor);
+    }
+
+    @Override
+    public int getControlBehavior() {
+        return RuleConstant.CONTROL_BEHAVIOR_WARM_UP;
+    }
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpRateLimiterControllerFactory.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.controller;
+
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingController;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory;
+
+import static com.alibaba.csp.sentinel.slots.block.flow.controller.WarmUpControllerFactory.coldFactor;
+
+public class WarmUpRateLimiterControllerFactory implements TrafficShapingControllerFactory {
+
+    @Override
+    public TrafficShapingController create(FlowRule rule) {
+        return new WarmUpRateLimiterController(rule.getCount(), rule.getWarmUpPeriodSec(), rule.getMaxQueueingTimeMs(), coldFactor);
+    }
+
+    @Override
+    public int getControlBehavior() {
+        return RuleConstant.CONTROL_BEHAVIOR_WARM_UP_RATE_LIMITER;
+    }
+}

sentinel-core/src/main/resources/META-INF/services/com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory

@@ -0,0 +1,4 @@
+com.alibaba.csp.sentinel.slots.block.flow.controller.DefaultControllerFactory
+com.alibaba.csp.sentinel.slots.block.flow.controller.WarmUpControllerFactory
+com.alibaba.csp.sentinel.slots.block.flow.controller.ThrottlingControllerFactory
+com.alibaba.csp.sentinel.slots.block.flow.controller.WarmUpRateLimiterControllerFactory