Refactor FlowRuleChecker to improve code reuse

Signed-off-by: Eric Zhao <sczyh16@gmail.com>

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) {
+    public void checkFlow(Function<String, Collection<FlowRule>> ruleProvider, ResourceWrapper resource,
-        return passCheck(rule, context, node, acquireCount, false);
+                          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,
-                                          boolean prioritized) {
+                                                    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,
+    private static boolean applyTokenResult(/*@NonNull*/ TokenResult result, FlowRule rule, Context context,
-                                                 int acquireCount, boolean prioritized) {
+                                                         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() {}
 }

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;
 
 /**
  * <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 {
+    void checkFlow(ResourceWrapper resource, Context context, DefaultNode node, int count, boolean prioritized)
-        // Flow rule map cannot be null.
+        throws BlockException {
-        Map<String, List<FlowRule>> flowRules = FlowRuleManager.getFlowRuleMap();
+        checker.checkFlow(ruleProvider, resource, context, node, count, prioritized);
-
-        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);
     }
 
     @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);
+        }
+    };
 }

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

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,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);