diff --git a/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowArgument.java b/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowArgument.java new file mode 100644 index 0000000000..575e6454c4 --- /dev/null +++ b/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowArgument.java @@ -0,0 +1,27 @@ +/* + * Copyright 1999-2018 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 + * + * http://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.param; + +/** + * ParamFlowArgument + */ +public interface ParamFlowArgument { + + /** + * @return the object as a key of param flow limit + */ + Object paramFlowKey(); +} diff --git a/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowChecker.java b/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowChecker.java index 055c0560a0..bd5ed8a0ea 100644 --- a/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowChecker.java +++ b/sentinel-extension/sentinel-parameter-flow-control/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowChecker.java @@ -56,8 +56,14 @@ public static boolean passCheck(ResourceWrapper resourceWrapper, /*@Valid*/ Para return true; } - // Get parameter value. If value is null, then pass. + // Get parameter value. Object value = args[paramIdx]; + + // Assign value with the result of paramFlowKey method + if (value instanceof ParamFlowArgument) { + value = ((ParamFlowArgument) value).paramFlowKey(); + } + // If value is null, then pass if (value == null) { return true; } diff --git a/sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowCheckerTest.java b/sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowCheckerTest.java index d7010b7ad1..bc545c6f84 100644 --- a/sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowCheckerTest.java +++ b/sentinel-extension/sentinel-parameter-flow-control/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/param/ParamFlowCheckerTest.java @@ -15,29 +15,27 @@ */ package com.alibaba.csp.sentinel.slots.block.flow.param; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; +import com.alibaba.csp.sentinel.EntryType; +import com.alibaba.csp.sentinel.slotchain.ResourceWrapper; +import com.alibaba.csp.sentinel.slotchain.StringResourceWrapper; +import com.alibaba.csp.sentinel.slots.block.RuleConstant; +import com.alibaba.csp.sentinel.slots.statistic.cache.ConcurrentLinkedHashMapWrapper; +import com.alibaba.csp.sentinel.util.TimeUtil; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; - -import com.alibaba.csp.sentinel.EntryType; -import com.alibaba.csp.sentinel.slotchain.ResourceWrapper; -import com.alibaba.csp.sentinel.slotchain.StringResourceWrapper; -import com.alibaba.csp.sentinel.slots.block.RuleConstant; -import com.alibaba.csp.sentinel.slots.statistic.cache.ConcurrentLinkedHashMapWrapper; -import com.alibaba.csp.sentinel.util.TimeUtil; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; /** * Test cases for {@link ParamFlowChecker}. @@ -58,7 +56,7 @@ public void testHotParamCheckerPassCheckExceedArgs() { rule.setParamIdx(paramIdx); assertTrue("The rule will pass if the paramIdx exceeds provided args", - ParamFlowChecker.passCheck(resourceWrapper, rule, 1, "abc")); + ParamFlowChecker.passCheck(resourceWrapper, rule, 1, "abc")); } @Test @@ -109,7 +107,7 @@ public void testSingleValueCheckThreadCountWithExceptionItems() { int thresholdD = 7; ParamFlowRule rule = new ParamFlowRule(resourceName).setCount(globalThreshold).setParamIdx(paramIdx) - .setGrade(RuleConstant.FLOW_GRADE_THREAD); + .setGrade(RuleConstant.FLOW_GRADE_THREAD); String valueA = "valueA"; String valueB = "valueB"; @@ -137,7 +135,7 @@ public void testSingleValueCheckThreadCountWithExceptionItems() { when(metric.getThreadCount(paramIdx, valueA)).thenReturn(globalThreshold); when(metric.getThreadCount(paramIdx, valueB)).thenReturn(thresholdB - 1L); when(metric.getThreadCount(paramIdx, valueC)).thenReturn(globalThreshold + 1); - when(metric.getThreadCount(paramIdx, valueD)).thenReturn(globalThreshold - 1).thenReturn((long)thresholdD); + when(metric.getThreadCount(paramIdx, valueD)).thenReturn(globalThreshold - 1).thenReturn((long) thresholdD); assertFalse(ParamFlowChecker.passSingleValueCheck(resourceWrapper, rule, 1, valueA)); assertTrue(ParamFlowChecker.passSingleValueCheck(resourceWrapper, rule, 1, valueB)); @@ -174,12 +172,12 @@ public void testPassLocalCheckForArray() throws InterruptedException { double globalThreshold = 1; ParamFlowRule rule = new ParamFlowRule(resourceName).setParamIdx(paramIdx) - .setControlBehavior(RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER).setCount(globalThreshold); + .setControlBehavior(RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER).setCount(globalThreshold); TimeUtil.currentTimeMillis(); String v1 = "a", v2 = "B", v3 = "Cc"; - Object arr = new String[] {v1, v2, v3}; + Object arr = new String[]{v1, v2, v3}; ParameterMetric metric = new ParameterMetric(); ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric); metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper(4000)); @@ -188,6 +186,41 @@ public void testPassLocalCheckForArray() throws InterruptedException { assertFalse(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, arr)); } + @Test + public void testPassLocalCheckForComplexParam() throws InterruptedException { + class User implements ParamFlowArgument { + Integer id; + String name; + String address; + + public User(Integer id, String name, String address) { + this.id = id; + this.name = name; + this.address = address; + } + + @Override + public Object paramFlowKey() { + return name; + } + } + final String resourceName = "testPassLocalCheckForComplexParam"; + final ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN); + int paramIdx = 0; + double globalThreshold = 1; + + ParamFlowRule rule = new ParamFlowRule(resourceName).setParamIdx(paramIdx).setCount(globalThreshold); + + Object[] args = new Object[]{new User(1, "Bob", "Hangzhou"), 10, "Demo"}; + ParameterMetric metric = new ParameterMetric(); + ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric); + metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper(4000)); + metric.getRuleTokenCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper(4000)); + + assertTrue(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, args)); + assertFalse(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, args)); + } + @Before public void setUp() throws Exception { ParameterMetricStorage.getMetricsMap().clear();