diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/RuleConstant.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/RuleConstant.java index 0b157675c5..c975e000b9 100755 --- a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/RuleConstant.java +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/RuleConstant.java @@ -50,6 +50,10 @@ public final class RuleConstant { public static final int CONTROL_BEHAVIOR_WARM_UP = 1; public static final int CONTROL_BEHAVIOR_RATE_LIMITER = 2; public static final int CONTROL_BEHAVIOR_WARM_UP_RATE_LIMITER = 3; + /** + * User-defined control behavior values MUST be >= 256 (0x100) + */ + public static final int CONTROL_BEHAVIOR_USER_DEFINED_MIN = 0x100; // 256 public static final int DEFAULT_BLOCK_STRATEGY = 0; public static final int TRY_AGAIN_BLOCK_STRATEGY = 1; diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleUtil.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleUtil.java index bc28a9dccb..03caee98a6 100644 --- a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleUtil.java +++ b/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 Map> buildFlowRuleMap(List 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()); diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/LoggingTrafficShapingControllerFactory.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/LoggingTrafficShapingControllerFactory.java new file mode 100644 index 0000000000..64e1bbc42e --- /dev/null +++ b/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(); + } +} diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactories.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactories.java new file mode 100644 index 0000000000..9b803fa4fe --- /dev/null +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactories.java @@ -0,0 +1,101 @@ +/* + * 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 com.alibaba.csp.sentinel.slots.block.RuleConstant; +import com.alibaba.csp.sentinel.spi.SpiLoader; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.BinaryOperator; +import java.util.stream.Collectors; + +public class TrafficShapingControllerFactories { + + private static final Map 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 usingExisting() { + return (existing, replacement) -> { + RecordLog.warn("[TrafficShapingControllerFactories] Duplicate control behavior [{}], " + + "using existing factory [{}], ignoring [{}]", + existing.getControlBehavior(), + existing.getClass().getName(), + replacement.getClass().getName()); + return existing; + }; + } + + private static TrafficShapingControllerFactory logging(TrafficShapingControllerFactory factory) { + return new LoggingTrafficShapingControllerFactory(factory); + } + + /** + * Validates the control behavior namespace of a factory. + */ + private static void validateControlBehavior(TrafficShapingControllerFactory factory) { + int controlBehavior = factory.getControlBehavior(); + + // User-defined factories must not use reserved range [0, 255] + if (!factory.isBuiltIn() && isReservedControlBehavior(controlBehavior)) { + throw new IllegalArgumentException(String.format( + "Invalid control behavior [%d] for factory [%s]. " + + "Control behavior values in range [0, %d] are reserved for Sentinel built-in implementations. " + + "User-defined factories must use values >= %d to ensure compatibility with future Sentinel upgrades.", + controlBehavior, + factory.getClass().getName(), + RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN - 1, + RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN)); + } + + RecordLog.info("[TrafficShapingControllerFactories] Registered factory [{}] for control behavior [{}]", + factory.getClass().getName(), controlBehavior); + } + + private static Map initFactories() { + List factories = SpiLoader.of(TrafficShapingControllerFactory.class) + .loadInstanceListSorted(); + // Validate all factories + for (TrafficShapingControllerFactory factory : factories) { + validateControlBehavior(factory); + } + return factories.stream() + .collect(Collectors.toMap( + TrafficShapingControllerFactory::getControlBehavior, + TrafficShapingControllerFactories::logging, + usingExisting(), + HashMap::new)); + } + + public static TrafficShapingControllerFactory get(int controlBehavior) { + return FACTORIES.get(controlBehavior); + } + + private TrafficShapingControllerFactories() { + } + + public static boolean isReservedControlBehavior(int controlBehavior) { + return controlBehavior < RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN; + } +} diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactory.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactory.java new file mode 100644 index 0000000000..5804d7f051 --- /dev/null +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactory.java @@ -0,0 +1,49 @@ +/* + * 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(); + + /** + * Indicates whether this factory is a built-in Sentinel implementation. + * Built-in factories are allowed to use control behavior values in the reserved range [0, 255]. + * User-defined factories should return {@code false} (default) and use values >= 256. + * + * This method is used internally for validation during factory registration to ensure + * proper namespace separation and prevent conflicts. + * + * @return {@code true} if this is a Sentinel built-in factory, {@code false} for user-defined implementations + */ + default boolean isBuiltIn() { + return false; + } +} diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/DefaultControllerFactory.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/DefaultControllerFactory.java new file mode 100644 index 0000000000..081104ffe5 --- /dev/null +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/DefaultControllerFactory.java @@ -0,0 +1,39 @@ +/* + * 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; + } + + @Override + public boolean isBuiltIn() { + return true; + } +} diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/ThrottlingControllerFactory.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/ThrottlingControllerFactory.java new file mode 100644 index 0000000000..60124bc550 --- /dev/null +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/ThrottlingControllerFactory.java @@ -0,0 +1,39 @@ +/* + * 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; + } + + @Override + public boolean isBuiltIn() { + return true; + } +} diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpControllerFactory.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpControllerFactory.java new file mode 100644 index 0000000000..af53c1443c --- /dev/null +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpControllerFactory.java @@ -0,0 +1,42 @@ +/* + * 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; + } + + @Override + public boolean isBuiltIn() { + return true; + } +} diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpRateLimiterControllerFactory.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpRateLimiterControllerFactory.java new file mode 100644 index 0000000000..9b37a4bd76 --- /dev/null +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/controller/WarmUpRateLimiterControllerFactory.java @@ -0,0 +1,41 @@ +/* + * 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; + } + + @Override + public boolean isBuiltIn() { + return true; + } +} diff --git a/sentinel-core/src/main/resources/META-INF/services/com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory b/sentinel-core/src/main/resources/META-INF/services/com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory new file mode 100644 index 0000000000..97579a6f1a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactoriesTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactoriesTest.java new file mode 100644 index 0000000000..53d0b011ff --- /dev/null +++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/TrafficShapingControllerFactoriesTest.java @@ -0,0 +1,107 @@ +package com.alibaba.csp.sentinel.slots.block.flow; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import com.alibaba.csp.sentinel.slots.block.RuleConstant; +import com.alibaba.csp.sentinel.slots.block.flow.controller.CustomTokenBucketControllerFactory; +import com.alibaba.csp.sentinel.slots.block.flow.controller.CustomTokenBucketControllerFactory.CustomTokenBucketController; +import java.util.ArrayList; +import java.util.List; +import org.junit.Test; + +/** + * Test cases for {@link TrafficShapingControllerFactories}. + * + * @author icodening + */ +public class TrafficShapingControllerFactoriesTest { + + + + @Test + public void testGetBuiltInFactories() { + // Test that all built-in factories are registered + assertNotNull("Default factory should be registered", + TrafficShapingControllerFactories.get(RuleConstant.CONTROL_BEHAVIOR_DEFAULT)); + assertNotNull("WarmUp factory should be registered", + TrafficShapingControllerFactories.get(RuleConstant.CONTROL_BEHAVIOR_WARM_UP)); + assertNotNull("RateLimiter factory should be registered", + TrafficShapingControllerFactories.get(RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER)); + assertNotNull("WarmUpRateLimiter factory should be registered", + TrafficShapingControllerFactories.get(RuleConstant.CONTROL_BEHAVIOR_WARM_UP_RATE_LIMITER)); + assertNotNull("CustomTokenBucket factory should be registered", + TrafficShapingControllerFactories.get(RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN)); + } + + @Test + public void testGetNonExistentFactory() { + // Test that non-existent control behavior returns null + assertNull("Non-existent factory should return null", + TrafficShapingControllerFactories.get(999)); + } + + + @Test + public void testDefaultFactoryCreatesController() { + FlowRule rule = new FlowRule("test") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS) + .setControlBehavior(RuleConstant.CONTROL_BEHAVIOR_DEFAULT); + + TrafficShapingControllerFactory factory = TrafficShapingControllerFactories.get( + RuleConstant.CONTROL_BEHAVIOR_DEFAULT); + assertNotNull("Factory should exist", factory); + + TrafficShapingController controller = factory.create(rule); + assertNotNull("Controller should be created", controller); + } + @Test + public void testCustomControllerIntegrationWithFlowRuleUtil() { + // Create a rule with custom control behavior + FlowRule rule = new FlowRule("integrationTestResource") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS) + .setControlBehavior(CustomTokenBucketControllerFactory.CUSTOM_CONTROL_BEHAVIOR); + + List rules = new ArrayList<>(); + rules.add(rule); + FlowRuleUtil.buildFlowRuleMap(rules, null, false); + + assertNotNull("Rule should have a controller", rule.getRater()); + } + + @Test + public void testCustomControllerWithFlowRule() { + FlowRule rule = new FlowRule("customResource") + .setCount(50) + .setGrade(RuleConstant.FLOW_GRADE_QPS) + .setControlBehavior(CustomTokenBucketControllerFactory.CUSTOM_CONTROL_BEHAVIOR); + + CustomTokenBucketControllerFactory factory = new CustomTokenBucketControllerFactory(); + TrafficShapingController controller = factory.create(rule); + + assertNotNull("Controller should be created from rule", controller); + assertTrue("Should create CustomTokenBucketController from rule", + controller instanceof CustomTokenBucketController); + + assertTrue("Should allow initial request", controller.canPass(null, 1)); + assertFalse("Third request should fail", controller.canPass(null, 2)); + } + + + @Test + public void testUserDefinedControlBehaviorNotInReservedRange() { + int userDefinedBehavior = RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN; + assertFalse("User-defined control behavior should not be in reserved range", + TrafficShapingControllerFactories.isReservedControlBehavior(userDefinedBehavior)); + + assertTrue("255 should be in reserved range", + TrafficShapingControllerFactories.isReservedControlBehavior(255)); + assertFalse("256 should not be in reserved range", + TrafficShapingControllerFactories.isReservedControlBehavior(256)); + } +} + diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/BuiltInControllerFactoryTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/BuiltInControllerFactoryTest.java new file mode 100644 index 0000000000..9d1e59c856 --- /dev/null +++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/BuiltInControllerFactoryTest.java @@ -0,0 +1,168 @@ + +package com.alibaba.csp.sentinel.slots.block.flow.controller; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +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.TrafficShapingControllerFactories; +import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory; +import com.alibaba.csp.sentinel.slots.block.flow.controller.CustomTokenBucketControllerFactory.CustomTokenBucketController; +import java.lang.reflect.Method; +import org.junit.Test; + +/** + * Test cases for built-in controller factories. + * + * @author soulx + */ +public class BuiltInControllerFactoryTest { + + @Test + public void testDefaultControllerFactory() { + TrafficShapingControllerFactory factory = new DefaultControllerFactory(); + + assertTrue("DefaultControllerFactory should be built-in", factory.isBuiltIn()); + assertEquals("Control behavior should be DEFAULT", + RuleConstant.CONTROL_BEHAVIOR_DEFAULT, factory.getControlBehavior()); + + FlowRule rule = new FlowRule("test") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS); + + TrafficShapingController controller = factory.create(rule); + assertNotNull("Controller should be created", controller); + assertTrue("Should create DefaultController", controller instanceof DefaultController); + } + + @Test + public void testWarmUpControllerFactory() { + TrafficShapingControllerFactory factory = new WarmUpControllerFactory(); + + assertTrue("WarmUpControllerFactory should be built-in", factory.isBuiltIn()); + assertEquals("Control behavior should be WARM_UP", + RuleConstant.CONTROL_BEHAVIOR_WARM_UP, factory.getControlBehavior()); + + FlowRule rule = new FlowRule("test") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS) + .setWarmUpPeriodSec(10); + + TrafficShapingController controller = factory.create(rule); + assertNotNull("Controller should be created", controller); + assertTrue("Should create WarmUpController", controller instanceof WarmUpController); + } + + @Test + public void testThrottlingControllerFactory() { + TrafficShapingControllerFactory factory = new ThrottlingControllerFactory(); + + assertTrue("ThrottlingControllerFactory should be built-in", factory.isBuiltIn()); + assertEquals("Control behavior should be RATE_LIMITER", + RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER, factory.getControlBehavior()); + + FlowRule rule = new FlowRule("test") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS) + .setMaxQueueingTimeMs(500); + + TrafficShapingController controller = factory.create(rule); + assertNotNull("Controller should be created", controller); + assertTrue("Should create ThrottlingController", controller instanceof ThrottlingController); + } + + @Test + public void testWarmUpRateLimiterControllerFactory() { + TrafficShapingControllerFactory factory = new WarmUpRateLimiterControllerFactory(); + + assertTrue("WarmUpRateLimiterControllerFactory should be built-in", factory.isBuiltIn()); + assertEquals("Control behavior should be WARM_UP_RATE_LIMITER", + RuleConstant.CONTROL_BEHAVIOR_WARM_UP_RATE_LIMITER, factory.getControlBehavior()); + + FlowRule rule = new FlowRule("test") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS) + .setWarmUpPeriodSec(10) + .setMaxQueueingTimeMs(500); + + TrafficShapingController controller = factory.create(rule); + assertNotNull("Controller should be created", controller); + assertTrue("Should create WarmUpRateLimiterController", + controller instanceof WarmUpRateLimiterController); + } + + @Test + public void testCustomTokenBucketControllerFactory() { + TrafficShapingControllerFactory factory = new CustomTokenBucketControllerFactory(); + + assertFalse("CustomTokenBucketControllerFactory should be built-in", factory.isBuiltIn()); + assertEquals("Control behavior should be USER_DEFINED_MIN", + RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN, factory.getControlBehavior()); + + FlowRule rule = new FlowRule("test") + .setCount(100) + .setGrade(RuleConstant.FLOW_GRADE_QPS); + + TrafficShapingController controller = factory.create(rule); + assertNotNull("Controller should be created", controller); + assertTrue("Should create CustomTokenBucketController", + controller instanceof CustomTokenBucketController); + } + + @Test + public void testInvalidReservedBehaviorFactory() { + InvalidReservedBehaviorFactory invalidFactory = new InvalidReservedBehaviorFactory(); + assertFalse("Test factory should not be built-in", invalidFactory.isBuiltIn()); + int controlBehavior = invalidFactory.getControlBehavior(); + + try { + Method validateMethod = TrafficShapingControllerFactories.class + .getDeclaredMethod("validateControlBehavior", TrafficShapingControllerFactory.class); + validateMethod.setAccessible(true); + + try { + validateMethod.invoke(null, invalidFactory); + } catch (Exception e) { + // Expected: InvocationTargetException wrapping IllegalArgumentException + Throwable cause = e.getCause(); + assertTrue("Should throw IllegalArgumentException", + cause instanceof IllegalArgumentException); + + String errorMessage = cause.getMessage(); + assertTrue("Error message should mention reserved range", + errorMessage.contains("reserved")); + assertTrue("Error message should mention control behavior value", + errorMessage.contains("[" + controlBehavior + "]")); + assertTrue("Error message should mention user-defined minimum", + errorMessage.contains(String.valueOf(RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN))); + } + } catch (NoSuchMethodException e) { + fail("validateControlBehavior method should exist: " + e.getMessage()); + } + } + + @Test + public void testAllBuiltInFactoriesUseReservedRange() { + // Test that all built-in factories use control behavior in reserved range [0, 255] + TrafficShapingControllerFactory[] factories = { + new DefaultControllerFactory(), + new WarmUpControllerFactory(), + new ThrottlingControllerFactory(), + new WarmUpRateLimiterControllerFactory() + }; + + for (TrafficShapingControllerFactory factory : factories) { + int controlBehavior = factory.getControlBehavior(); + assertTrue( + String.format("Built-in factory %s should use reserved control behavior, but got %d", + factory.getClass().getSimpleName(), controlBehavior), + TrafficShapingControllerFactories.isReservedControlBehavior(controlBehavior)); + } + } +} + diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/CustomTokenBucketControllerFactory.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/CustomTokenBucketControllerFactory.java new file mode 100644 index 0000000000..e571f583b1 --- /dev/null +++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/CustomTokenBucketControllerFactory.java @@ -0,0 +1,40 @@ +package com.alibaba.csp.sentinel.slots.block.flow.controller; + +import com.alibaba.csp.sentinel.node.Node; +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 CustomTokenBucketControllerFactory implements TrafficShapingControllerFactory { + /** + * Custom control behavior constant - must be >= 256 + */ + public static final int CUSTOM_CONTROL_BEHAVIOR = RuleConstant.CONTROL_BEHAVIOR_USER_DEFINED_MIN; + @Override + public TrafficShapingController create(FlowRule rule) { + return new CustomTokenBucketController(); + } + + @Override + public int getControlBehavior() { + return CUSTOM_CONTROL_BEHAVIOR; + } + + public static class CustomTokenBucketController implements TrafficShapingController { + + + @Override + public boolean canPass(Node node, int acquireCount) { + return canPass(node, acquireCount, false); + } + + @Override + public boolean canPass(Node node, int acquireCount, boolean prioritized) { + if (acquireCount <= 1) return true; + return false; + } + + + } +} \ No newline at end of file diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/InvalidReservedBehaviorFactory.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/InvalidReservedBehaviorFactory.java new file mode 100644 index 0000000000..33277019ba --- /dev/null +++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/InvalidReservedBehaviorFactory.java @@ -0,0 +1,24 @@ +package com.alibaba.csp.sentinel.slots.block.flow.controller; + +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 com.alibaba.csp.sentinel.slots.block.flow.controller.CustomTokenBucketControllerFactory.CustomTokenBucketController; + +public class InvalidReservedBehaviorFactory implements TrafficShapingControllerFactory { + + @Override + public TrafficShapingController create(FlowRule rule) { + return new CustomTokenBucketController(); + } + + @Override + public int getControlBehavior() { + return 5; // This is in the reserved range [0, 255] + } + + @Override + public boolean isBuiltIn() { + return false; // Not a built-in factory + } + } \ No newline at end of file diff --git a/sentinel-core/src/test/resources/META-INF/services/com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory b/sentinel-core/src/test/resources/META-INF/services/com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory new file mode 100644 index 0000000000..35198fa61b --- /dev/null +++ b/sentinel-core/src/test/resources/META-INF/services/com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingControllerFactory @@ -0,0 +1,2 @@ +com.alibaba.csp.sentinel.slots.block.flow.controller.CustomTokenBucketControllerFactory +#com.alibaba.csp.sentinel.slots.block.flow.controller.InvalidReservedBehaviorFactory \ No newline at end of file