diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleChecker.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleChecker.java index 4799bfb3..26d34d39 100644 --- a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleChecker.java +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleChecker.java @@ -15,6 +15,8 @@ */ package com.alibaba.csp.sentinel.slots.block.flow; +import java.util.Collection; + import com.alibaba.csp.sentinel.cluster.ClusterStateManager; import com.alibaba.csp.sentinel.cluster.server.EmbeddedClusterTokenServerProvider; import com.alibaba.csp.sentinel.cluster.client.TokenClientProvider; @@ -25,23 +27,42 @@ import com.alibaba.csp.sentinel.context.Context; import com.alibaba.csp.sentinel.log.RecordLog; import com.alibaba.csp.sentinel.node.DefaultNode; import com.alibaba.csp.sentinel.node.Node; +import com.alibaba.csp.sentinel.slotchain.ResourceWrapper; +import com.alibaba.csp.sentinel.slots.block.BlockException; import com.alibaba.csp.sentinel.slots.block.RuleConstant; import com.alibaba.csp.sentinel.slots.clusterbuilder.ClusterBuilderSlot; import com.alibaba.csp.sentinel.util.StringUtil; +import com.alibaba.csp.sentinel.util.function.Function; /** * Rule checker for flow control rules. * * @author Eric Zhao */ -final class FlowRuleChecker { +public class FlowRuleChecker { - static boolean passCheck(/*@NonNull*/ FlowRule rule, Context context, DefaultNode node, int acquireCount) { - return passCheck(rule, context, node, acquireCount, false); + public void checkFlow(Function> ruleProvider, ResourceWrapper resource, + Context context, DefaultNode node, int count, boolean prioritized) throws BlockException { + if (ruleProvider == null || resource == null) { + return; + } + Collection rules = ruleProvider.apply(resource.getName()); + if (rules != null) { + for (FlowRule rule : rules) { + if (!canPassCheck(rule, context, node, count, prioritized)) { + throw new FlowException(rule.getLimitApp(), rule); + } + } + } } - static boolean passCheck(/*@NonNull*/ FlowRule rule, Context context, DefaultNode node, int acquireCount, - boolean prioritized) { + public boolean canPassCheck(/*@NonNull*/ FlowRule rule, Context context, DefaultNode node, + int acquireCount) { + return canPassCheck(rule, context, node, acquireCount, false); + } + + public boolean canPassCheck(/*@NonNull*/ FlowRule rule, Context context, DefaultNode node, int acquireCount, + boolean prioritized) { String limitApp = rule.getLimitApp(); if (limitApp == null) { return true; @@ -162,8 +183,9 @@ final class FlowRuleChecker { return null; } - private static boolean applyTokenResult(/*@NonNull*/ TokenResult result, FlowRule rule, Context context, DefaultNode node, - int acquireCount, boolean prioritized) { + private static boolean applyTokenResult(/*@NonNull*/ TokenResult result, FlowRule rule, Context context, + DefaultNode node, + int acquireCount, boolean prioritized) { switch (result.getStatus()) { case TokenResultStatus.OK: return true; @@ -185,6 +207,4 @@ final class FlowRuleChecker { return false; } } - - private FlowRuleChecker() {} } \ No newline at end of file diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlot.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlot.java index 9aa900ec..f99987f7 100755 --- a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlot.java +++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlot.java @@ -15,6 +15,7 @@ */ package com.alibaba.csp.sentinel.slots.block.flow; +import java.util.Collection; import java.util.List; import java.util.Map; @@ -23,6 +24,8 @@ import com.alibaba.csp.sentinel.node.DefaultNode; import com.alibaba.csp.sentinel.slotchain.AbstractLinkedProcessorSlot; import com.alibaba.csp.sentinel.slotchain.ResourceWrapper; import com.alibaba.csp.sentinel.slots.block.BlockException; +import com.alibaba.csp.sentinel.util.AssertUtil; +import com.alibaba.csp.sentinel.util.function.Function; /** *

@@ -135,6 +138,23 @@ import com.alibaba.csp.sentinel.slots.block.BlockException; */ public class FlowSlot extends AbstractLinkedProcessorSlot { + private final FlowRuleChecker checker; + + public FlowSlot() { + this(new FlowRuleChecker()); + } + + /** + * Package-private for test. + * + * @param checker flow rule checker + * @since 1.6.1 + */ + FlowSlot(FlowRuleChecker checker) { + AssertUtil.notNull(checker, "flow checker should not be null"); + this.checker = checker; + } + @Override public void entry(Context context, ResourceWrapper resourceWrapper, DefaultNode node, int count, boolean prioritized, Object... args) throws Throwable { @@ -143,26 +163,22 @@ public class FlowSlot extends AbstractLinkedProcessorSlot { fireEntry(context, resourceWrapper, node, count, prioritized, args); } - void checkFlow(ResourceWrapper resource, Context context, DefaultNode node, int count, boolean prioritized) throws BlockException { - // Flow rule map cannot be null. - Map> flowRules = FlowRuleManager.getFlowRuleMap(); - - List rules = flowRules.get(resource.getName()); - if (rules != null) { - for (FlowRule rule : rules) { - if (!canPassCheck(rule, context, node, count, prioritized)) { - throw new FlowException(rule.getLimitApp(), rule); - } - } - } - } - - boolean canPassCheck(FlowRule rule, Context context, DefaultNode node, int count, boolean prioritized) { - return FlowRuleChecker.passCheck(rule, context, node, count, prioritized); + void checkFlow(ResourceWrapper resource, Context context, DefaultNode node, int count, boolean prioritized) + throws BlockException { + checker.checkFlow(ruleProvider, resource, context, node, count, prioritized); } @Override public void exit(Context context, ResourceWrapper resourceWrapper, int count, Object... args) { fireExit(context, resourceWrapper, count, args); } + + private final Function> ruleProvider = new Function>() { + @Override + public Collection apply(String resource) { + // Flow rule map should not be null. + Map> flowRules = FlowRuleManager.getFlowRuleMap(); + return flowRules.get(resource); + } + }; } diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleCheckerTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleCheckerTest.java index 10c460f3..bdd8072d 100644 --- a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleCheckerTest.java +++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleCheckerTest.java @@ -149,7 +149,8 @@ public class FlowRuleCheckerTest { public void testPassCheckNullLimitApp() { FlowRule rule = new FlowRule("abc").setCount(1); rule.setLimitApp(null); - assertTrue(FlowRuleChecker.passCheck(rule, null, null, 1)); + FlowRuleChecker checker = new FlowRuleChecker(); + assertTrue(checker.canPassCheck(rule, null, null, 1)); } @Test @@ -161,7 +162,8 @@ public class FlowRuleCheckerTest { Context context = mock(Context.class); when(context.getOrigin()).thenReturn("def"); - assertTrue(FlowRuleChecker.passCheck(rule, context, node, 1)); + FlowRuleChecker checker = new FlowRuleChecker(); + assertTrue(checker.canPassCheck(rule, context, node, 1)); } @Before diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlotTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlotTest.java index ee5de0a8..66026e2b 100644 --- a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlotTest.java +++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlotTest.java @@ -23,6 +23,7 @@ import com.alibaba.csp.sentinel.context.ContextTestUtil; import com.alibaba.csp.sentinel.node.DefaultNode; import com.alibaba.csp.sentinel.slotchain.ResourceWrapper; import com.alibaba.csp.sentinel.slotchain.StringResourceWrapper; +import com.alibaba.csp.sentinel.util.function.Function; import org.junit.After; import org.junit.Before; @@ -49,11 +50,13 @@ public class FlowSlotTest { } @Test + @SuppressWarnings("unchecked") public void testCheckFlowPass() throws Exception { - FlowSlot flowSlot = mock(FlowSlot.class); + FlowRuleChecker checker = mock(FlowRuleChecker.class); + FlowSlot flowSlot = new FlowSlot(checker); Context context = mock(Context.class); DefaultNode node = mock(DefaultNode.class); - doCallRealMethod().when(flowSlot).checkFlow(any(ResourceWrapper.class), any(Context.class), + doCallRealMethod().when(checker).checkFlow(any(Function.class), any(ResourceWrapper.class), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean()); String resA = "resAK"; @@ -63,9 +66,9 @@ public class FlowSlotTest { // Here we only load rules for resA. FlowRuleManager.loadRules(Collections.singletonList(rule1)); - when(flowSlot.canPassCheck(eq(rule1), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean())) + when(checker.canPassCheck(eq(rule1), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean())) .thenReturn(true); - when(flowSlot.canPassCheck(eq(rule2), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean())) + when(checker.canPassCheck(eq(rule2), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean())) .thenReturn(false); flowSlot.checkFlow(new StringResourceWrapper(resA, EntryType.IN), context, node, 1, false); @@ -73,20 +76,22 @@ public class FlowSlotTest { } @Test(expected = FlowException.class) + @SuppressWarnings("unchecked") public void testCheckFlowBlock() throws Exception { - FlowSlot flowSlot = mock(FlowSlot.class); + FlowRuleChecker checker = mock(FlowRuleChecker.class); + FlowSlot flowSlot = new FlowSlot(checker); Context context = mock(Context.class); DefaultNode node = mock(DefaultNode.class); - doCallRealMethod().when(flowSlot).checkFlow(any(ResourceWrapper.class), any(Context.class), + doCallRealMethod().when(checker).checkFlow(any(Function.class), any(ResourceWrapper.class), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean()); String resA = "resAK"; FlowRule rule = new FlowRule(resA).setCount(10); FlowRuleManager.loadRules(Collections.singletonList(rule)); - when(flowSlot.canPassCheck(any(FlowRule.class), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean())) + when(checker.canPassCheck(any(FlowRule.class), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean())) .thenReturn(false); flowSlot.checkFlow(new StringResourceWrapper(resA, EntryType.IN), context, node, 1, false); } -} \ No newline at end of file +}