crabapples
/
Sentinel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor FlowRuleChecker to improve code reuse 

Signed-off-by: Eric Zhao <sczyh16@gmail.com>
This commit is contained in:
Eric Zhao 2019-05-11 22:15:09 +08:00
parent 4ec0462e31
commit 54da16d304
4 changed files with 78 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleChecker.java
View File

@ -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,22 +27,41 @@ 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<String, Collection<FlowRule>> ruleProvider, ResourceWrapper resource,
Context context, DefaultNode node, int count, boolean prioritized) throws BlockException {
if (ruleProvider == null || resource == null) {
return;
}
Collection<FlowRule> 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,
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) {
@ -162,7 +183,8 @@ final class FlowRuleChecker {
return null;
}
private static boolean applyTokenResult(/*@NonNull*/ TokenResult result, FlowRule rule, Context context, DefaultNode node,
private static boolean applyTokenResult(/*@NonNull*/ TokenResult result, FlowRule rule, Context context,
DefaultNode node,
int acquireCount, boolean prioritized) {
switch (result.getStatus()) {
case TokenResultStatus.OK:
@ -185,6 +207,4 @@ final class FlowRuleChecker {
return false;
}
}
private FlowRuleChecker() {}
}

48
sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlot.java
View File

@ -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;
/**
* <p>
@ -135,6 +138,23 @@ import com.alibaba.csp.sentinel.slots.block.BlockException;
*/
public class FlowSlot extends AbstractLinkedProcessorSlot<DefaultNode> {
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<DefaultNode> {
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<String, List<FlowRule>> flowRules = FlowRuleManager.getFlowRuleMap();
List<FlowRule> 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<String, Collection<FlowRule>> ruleProvider = new Function<String, Collection<FlowRule>>() {
@Override
public Collection<FlowRule> apply(String resource) {
// Flow rule map should not be null.
Map<String, List<FlowRule>> flowRules = FlowRuleManager.getFlowRuleMap();
return flowRules.get(resource);
}
};
}

6
sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleCheckerTest.java
View File

@ -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

19
sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlotTest.java
View File

@ -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,18 +76,20 @@ 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);