com.alibaba.csp.sentinel.slotchain.ResourceWrapper Java Examples

private void checkGatewayParamFlow(ResourceWrapper resourceWrapper, int count, Object... args)
    throws BlockException {
    if (args == null) {
        return;
    }

    List<ParamFlowRule> rules = GatewayRuleManager.getConvertedParamRules(resourceWrapper.getName());
    if (rules == null || rules.isEmpty()) {
        return;
    }

    for (ParamFlowRule rule : rules) {
        // Initialize the parameter metrics.
        ParameterMetricStorage.initParamMetricsFor(resourceWrapper, rule);

        if (!ParamFlowChecker.passCheck(resourceWrapper, rule, count, args)) {
            String triggeredParam = "";
            if (args.length > rule.getParamIdx()) {
                Object value = args[rule.getParamIdx()];
                triggeredParam = String.valueOf(value);
            }
            throw new ParamFlowException(resourceWrapper.getName(), triggeredParam, rule);
        }
    }
}
 

private void checkGatewayParamFlow(ResourceWrapper resourceWrapper, int count, Object... args)
    throws BlockException {
    if (args == null) {
        return;
    }

    List<ParamFlowRule> rules = GatewayRuleManager.getConvertedParamRules(resourceWrapper.getName());
    if (rules == null || rules.isEmpty()) {
        return;
    }

    for (ParamFlowRule rule : rules) {
        // Initialize the parameter metrics.
        ParameterMetricStorage.initParamMetricsFor(resourceWrapper, rule);

        if (!ParamFlowChecker.passCheck(resourceWrapper, rule, count, args)) {
            String triggeredParam = "";
            if (args.length > rule.getParamIdx()) {
                Object value = args[rule.getParamIdx()];
                triggeredParam = String.valueOf(value);
            }
            throw new ParamFlowException(resourceWrapper.getName(), triggeredParam, rule);
        }
    }
}
 

@Test
public void testPassLocalCheckForArray() throws InterruptedException {
    final String resourceName = "testPassLocalCheckForArray";
    final ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    int paramIdx = 0;
    double globalThreshold = 1;

    ParamFlowRule rule = new ParamFlowRule(resourceName).setParamIdx(paramIdx)
        .setControlBehavior(RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER).setCount(globalThreshold);

    TimeUtil.currentTimeMillis();

    String v1 = "a", v2 = "B", v3 = "Cc";
    Object arr = new String[] {v1, v2, v3};
    ParameterMetric metric = new ParameterMetric();
    ParameterMetricStorage.getMetricsMap().put(resourceWrapper.getName(), metric);
    metric.getRuleTimeCounterMap().put(rule, new ConcurrentLinkedHashMapWrapper<Object, AtomicLong>(4000));

    assertTrue(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, arr));
    assertFalse(ParamFlowChecker.passCheck(resourceWrapper, rule, 1, arr));
}
 

void checkFlow(ResourceWrapper resourceWrapper, int count, Object... args) throws BlockException {
    if (args == null) {
        return;
    }
    if (!ParamFlowRuleManager.hasRules(resourceWrapper.getName())) {
        return;
    }
    List<ParamFlowRule> rules = ParamFlowRuleManager.getRulesOfResource(resourceWrapper.getName());

    for (ParamFlowRule rule : rules) {
        applyRealParamIdx(rule, args.length);

        // Initialize the parameter metrics.
        ParameterMetricStorage.initParamMetricsFor(resourceWrapper, rule);

        if (!ParamFlowChecker.passCheck(resourceWrapper, rule, count, args)) {
            String triggeredParam = "";
            if (args.length > rule.getParamIdx()) {
                Object value = args[rule.getParamIdx()];
                triggeredParam = String.valueOf(value);
            }
            throw new ParamFlowException(resourceWrapper.getName(), triggeredParam, rule);
        }
    }
}
 

@Test
@SuppressWarnings("unchecked")
public void testCheckFlowPass() throws Exception {
    FlowRuleChecker checker = mock(FlowRuleChecker.class);
    FlowSlot flowSlot = new FlowSlot(checker);
    Context context = mock(Context.class);
    DefaultNode node = mock(DefaultNode.class);
    doCallRealMethod().when(checker).checkFlow(any(Function.class), any(ResourceWrapper.class), any(Context.class),
        any(DefaultNode.class), anyInt(), anyBoolean());

    String resA = "resAK";
    String resB = "resBK";
    FlowRule rule1 = new FlowRule(resA).setCount(10);
    FlowRule rule2 = new FlowRule(resB).setCount(10);
    // Here we only load rules for resA.
    FlowRuleManager.loadRules(Collections.singletonList(rule1));

    when(checker.canPassCheck(eq(rule1), any(Context.class), any(DefaultNode.class), anyInt(), anyBoolean()))
        .thenReturn(true);
    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);
    flowSlot.checkFlow(new StringResourceWrapper(resB, EntryType.IN), context, node, 1, false);
}
 

private void testCustomContextEntryWithFullContextSize(String resourceName, boolean async) {
    fillFullContext();
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    String contextName = "custom-context-" + System.currentTimeMillis();
    ContextUtil.enter(contextName, "9527");

    // Prepare a slot that "should not pass". If entered the slot, exception will be thrown.
    addShouldNotPassSlotFor(resourceWrapper);

    Entry entry = null;
    try {
        if (async) {
            entry = ctSph.asyncEntry(resourceName, resourceWrapper.getEntryType(), 1);
        } else {
            entry = ctSph.entry(resourceWrapper, 1);
        }
    } catch (BlockException ex) {
        fail("Unexpected blocked: " + ex.getClass().getCanonicalName());
    } finally {
        if (entry != null) {
            entry.exit();
        }
        ContextUtil.exit();
    }
}
 

/**
 * Init the parameter metric and index map for given resource.
 * Package-private for test.
 *
 * @param resourceWrapper resource to init
 * @param rule            relevant rule
 */
public static void initParamMetricsFor(ResourceWrapper resourceWrapper, /*@Valid*/ ParamFlowRule rule) {
    if (resourceWrapper == null || resourceWrapper.getName() == null) {
        return;
    }
    String resourceName = resourceWrapper.getName();
    ParameterMetric metric;
    // Assume that the resource is valid.
    if ((metric = metricsMap.get(resourceName)) == null) {
        synchronized (LOCK) {
            if ((metric = metricsMap.get(resourceName)) == null) {
                metric = new ParameterMetric();
                metricsMap.put(resourceWrapper.getName(), metric);
                RecordLog.info("[ParameterMetricStorage] Creating parameter metric for: " + resourceWrapper.getName());
            }
        }
    }
    metric.initialize(rule);
}
 

@Test
public void testInitParamMetrics() {
    ParamFlowRule rule = new ParamFlowRule();
    rule.setParamIdx(1);
    int index = 1;
    String resourceName = "res-" + System.currentTimeMillis();
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);

    assertNull(ParameterMetricStorage.getParamMetric(resourceWrapper));

    ParameterMetricStorage.initParamMetricsFor(resourceWrapper, rule);
    ParameterMetric metric = ParameterMetricStorage.getParamMetric(resourceWrapper);
    assertNotNull(metric);
    assertNotNull(metric.getRuleTimeCounterMap().get(rule));
    assertNotNull(metric.getThreadCountMap().get(index));

    // Duplicate init.
    ParameterMetricStorage.initParamMetricsFor(resourceWrapper, rule);
    assertSame(metric, ParameterMetricStorage.getParamMetric(resourceWrapper));

    ParamFlowRule rule2 = new ParamFlowRule();
    rule2.setParamIdx(1);
    assertSame(metric, ParameterMetricStorage.getParamMetric(resourceWrapper));
}
 

@Test(expected = AuthorityException.class)
public void testCheckAuthorityNoExceptionItemsWhiteFail() throws Exception {
    String origin = "appA";
    String resourceName = "testCheckAuthorityNoExceptionItemsWhiteFail";
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    ContextUtil.enter("entrance", origin);
    try {
        AuthorityRule ruleB = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp("appB, appE")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_WHITE);

        AuthorityRuleManager.loadRules(Collections.singletonList(ruleB));
        authoritySlot.checkBlackWhiteAuthority(resourceWrapper, ContextUtil.getContext());
    } finally {
        ContextUtil.exit();
    }
}
 

void checkFlow(ResourceWrapper resourceWrapper, int count, Object... args) throws BlockException {
    if (args == null) {
        return;
    }
    if (!ParamFlowRuleManager.hasRules(resourceWrapper.getName())) {
        return;
    }
    List<ParamFlowRule> rules = ParamFlowRuleManager.getRulesOfResource(resourceWrapper.getName());

    for (ParamFlowRule rule : rules) {
        applyRealParamIdx(rule, args.length);

        // Initialize the parameter metrics.
        ParameterMetricStorage.initParamMetricsFor(resourceWrapper, rule);

        if (!ParamFlowChecker.passCheck(resourceWrapper, rule, count, args)) {
            String triggeredParam = "";
            if (args.length > rule.getParamIdx()) {
                Object value = args[rule.getParamIdx()];
                triggeredParam = String.valueOf(value);
            }
            throw new ParamFlowException(resourceWrapper.getName(), triggeredParam, rule);
        }
    }
}
 

private static boolean passClusterCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                        Object value) {
    try {
        Collection<Object> params = toCollection(value);

        TokenService clusterService = pickClusterService();
        if (clusterService == null) {
            // No available cluster client or server, fallback to local or
            // pass in need.
            return fallbackToLocalOrPass(resourceWrapper, rule, count, params);
        }

        TokenResult result = clusterService.requestParamToken(rule.getClusterConfig().getFlowId(), count, params);
        switch (result.getStatus()) {
            case TokenResultStatus.OK:
                return true;
            case TokenResultStatus.BLOCKED:
                return false;
            default:
                return fallbackToLocalOrPass(resourceWrapper, rule, count, params);
        }
    } catch (Throwable ex) {
        RecordLog.warn("[ParamFlowChecker] Request cluster token for parameter unexpected failed", ex);
        return fallbackToLocalOrPass(resourceWrapper, rule, count, value);
    }
}
 

static boolean passSingleValueCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                    Object value) {
    if (rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {
        if (rule.getControlBehavior() == RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER) {
            return passThrottleLocalCheck(resourceWrapper, rule, acquireCount, value);
        } else {
            return passDefaultLocalCheck(resourceWrapper, rule, acquireCount, value);
        }
    } else if (rule.getGrade() == RuleConstant.FLOW_GRADE_THREAD) {
        Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
        long threadCount = getParameterMetric(resourceWrapper).getThreadCount(rule.getParamIdx(), value);
        if (exclusionItems.contains(value)) {
            int itemThreshold = rule.getParsedHotItems().get(value);
            return ++threadCount <= itemThreshold;
        }
        long threshold = (long)rule.getCount();
        return ++threadCount <= threshold;
    }

    return true;
}
 

@Override
public void entry(Context context, ResourceWrapper resourceWrapper, DefaultNode node, int count,
                  boolean prioritized, Object... args) throws Throwable {
    if (!ParamFlowRuleManager.hasRules(resourceWrapper.getName())) {
        fireEntry(context, resourceWrapper, node, count, prioritized, args);
        return;
    }

    checkFlow(resourceWrapper, count, args);
    fireEntry(context, resourceWrapper, node, count, prioritized, args);
}
 

public static boolean passCheck(ResourceWrapper resourceWrapper, /*@Valid*/ ParamFlowRule rule, /*@Valid*/ int count,
                         Object... args) {
    if (args == null) {
        return true;
    }

    int paramIdx = rule.getParamIdx();
    if (args.length <= paramIdx) {
        return true;
    }

    // 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;
    }

    if (rule.isClusterMode() && rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {
        return passClusterCheck(resourceWrapper, rule, count, value);
    }

    return passLocalCheck(resourceWrapper, rule, count, value);
}
 

@Override
public void onBlocked(BlockException ex, Context context, ResourceWrapper resourceWrapper,
                      DefaultNode param, int count, Object... args) {
    for (MetricExtension m : MetricExtensionProvider.getMetricExtensions()) {
        m.addBlock(resourceWrapper.getName(), count, context.getOrigin(), ex, args);
    }
}
 

@Test
public void testAsyncEntryNestedInSyncEntryNormalBlocked() {
    String previousResourceName = "fff";
    String resourceName = "testAsyncEntryNestedInSyncEntryNormalBlocked";
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);

    // Prepare a slot that "must block".
    MustBlockSlot slot = addMustBlockSlot(resourceWrapper);
    assertFalse(slot.exited);
    // Previous entry should pass.
    addShouldPassSlotFor(new StringResourceWrapper(previousResourceName, EntryType.IN));
    ContextUtil.enter("bcd-" + System.currentTimeMillis());

    AsyncEntry entry = null;
    Entry syncEntry = null;
    Entry previousEntry = null;
    try {
        // First enter a sync resource.
        syncEntry = ctSph.entry(previousResourceName, EntryType.IN, 1);
        // Record current entry (previous for next).
        previousEntry = ContextUtil.getContext().getCurEntry();
        // Then enter an async resource.
        entry = ctSph.asyncEntry(resourceName, EntryType.IN, 1);

        // Should not pass here.
    } catch (BlockException ex) {
        assertNotNull(previousEntry);
        assertNull(entry);
        assertTrue(slot.exited);
        assertSame(previousEntry, ContextUtil.getContext().getCurEntry());
        return;
    } finally {
        assertNull(entry);
        assertNotNull(syncEntry);

        syncEntry.exit();
        ContextUtil.exit();
    }
    fail("This async entry is expected to be blocked");
}
 

@Override
public void entry(Context context, ResourceWrapper resource, DefaultNode node, int count,
                  boolean prioritized, Object... args) throws Throwable {
    checkGatewayParamFlow(resource, count, args);

    fireEntry(context, resource, node, count, prioritized, args);
}
 

private void testDefaultContextEntryWithFullContextSize(String resourceName, boolean async) {
    fillFullContext();
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);

    // Prepare a slot that "should pass".
    ShouldPassSlot slot = addShouldPassSlotFor(resourceWrapper);
    assertFalse(slot.entered || slot.exited);

    Entry entry = null;
    try {
        if (!async) {
            entry = ctSph.entry(resourceWrapper, 1);
        } else {
            entry = ctSph.asyncEntry(resourceName, resourceWrapper.getEntryType(), 1);
            Context asyncContext = ((AsyncEntry)entry).getAsyncContext();
            assertTrue(ContextUtil.isDefaultContext(asyncContext));
            assertTrue(asyncContext.isAsync());
        }
        assertTrue(ContextUtil.isDefaultContext(ContextUtil.getContext()));
        assertTrue(slot.entered);
    } catch (BlockException ex) {
        fail("Unexpected blocked: " + ex.getClass().getCanonicalName());
    } finally {
        if (entry != null) {
            entry.exit();
            assertTrue(slot.exited);
        }
    }
}
 

private void testDefaultContextEntryWithFullContextSize(String resourceName, boolean async) {
    fillFullContext();
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);

    // Prepare a slot that "should pass".
    ShouldPassSlot slot = addShouldPassSlotFor(resourceWrapper);
    assertFalse(slot.entered || slot.exited);

    Entry entry = null;
    try {
        if (!async) {
            entry = ctSph.entry(resourceWrapper, 1);
        } else {
            entry = ctSph.asyncEntry(resourceName, resourceWrapper.getType(), 1);
            Context asyncContext = ((AsyncEntry)entry).getAsyncContext();
            assertTrue(ContextUtil.isDefaultContext(asyncContext));
            assertTrue(asyncContext.isAsync());
        }
        assertTrue(ContextUtil.isDefaultContext(ContextUtil.getContext()));
        assertTrue(slot.entered);
    } catch (BlockException ex) {
        fail("Unexpected blocked: " + ex.getClass().getCanonicalName());
    } finally {
        if (entry != null) {
            entry.exit();
            assertTrue(slot.exited);
        }
    }
}
 

@Test
public void testHotParamCheckerPassCheckExceedArgs() {
    final String resourceName = "testHotParamCheckerPassCheckExceedArgs";
    final ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    int paramIdx = 1;

    ParamFlowRule rule = new ParamFlowRule();
    rule.setResource(resourceName);
    rule.setCount(10);
    rule.setParamIdx(paramIdx);

    assertTrue("The rule will pass if the paramIdx exceeds provided args",
        ParamFlowChecker.passCheck(resourceWrapper, rule, 1, "abc"));
}
 

@Test
public void testHotParamCheckerPassCheckExceedArgs() {
    final String resourceName = "testHotParamCheckerPassCheckExceedArgs";
    final ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    int paramIdx = 1;

    ParamFlowRule rule = new ParamFlowRule();
    rule.setResource(resourceName);
    rule.setCount(10);
    rule.setParamIdx(paramIdx);

    assertTrue("The rule will pass if the paramIdx exceeds provided args",
            ParamFlowChecker.passCheck(resourceWrapper, rule, 1, "abc"));
}
 

@Override
public void onExit(Context context, ResourceWrapper resourceWrapper, int count, Object... args) {
    if (context.getCurEntry().getBlockError() == null) {
        ParameterMetric parameterMetric = ParameterMetricStorage.getParamMetric(resourceWrapper);

        if (parameterMetric != null) {
            parameterMetric.decreaseThreadCount(args);
        }
    }
}
 

@Test
public void testEntryAndAsyncEntryWhenSwitchOff() {
    // Turn off the switch.
    Constants.ON = false;

    String resourceNameA = "resSync";
    String resourceNameB = "resAsync";
    ResourceWrapper resourceWrapperA = new StringResourceWrapper(resourceNameA, EntryType.IN);
    ResourceWrapper resourceWrapperB = new StringResourceWrapper(resourceNameB, EntryType.IN);

    // Prepare a slot that "should not pass". If entered the slot, exception will be thrown.
    addShouldNotPassSlotFor(resourceWrapperA);
    addShouldNotPassSlotFor(resourceWrapperB);

    Entry entry = null;
    AsyncEntry asyncEntry = null;
    try {
        entry = ctSph.entry(resourceWrapperA, 1);
        asyncEntry = ctSph.asyncEntry(resourceNameB, resourceWrapperB.getEntryType(), 1);
    } catch (BlockException ex) {
        fail("Unexpected blocked: " + ex.getClass().getCanonicalName());
    } finally {
        if (asyncEntry != null) {
            asyncEntry.exit();
        }
        if (entry != null) {
            entry.exit();
        }
        Constants.ON = true;
    }
}
 

@Test
public void testLookUpSlotChain() {
    ResourceWrapper r1 = new StringResourceWrapper("firstRes", EntryType.IN);
    assertFalse(CtSph.getChainMap().containsKey(r1));
    ProcessorSlot<Object> chainR1 = ctSph.lookProcessChain(r1);
    assertNotNull("The slot chain for r1 should be created", chainR1);
    assertSame("Should return the cached slot chain once it has been created", chainR1, ctSph.lookProcessChain(r1));

    fillFullResources();
    ResourceWrapper r2 = new StringResourceWrapper("secondRes", EntryType.IN);
    assertFalse(CtSph.getChainMap().containsKey(r2));
    assertNull("The slot chain for r2 should not be created because amount exceeded", ctSph.lookProcessChain(r2));
    assertNull(ctSph.lookProcessChain(r2));
}
 

public static void checkDegrade(ResourceWrapper resource, Context context, DefaultNode node, int count)
    throws BlockException {

    Set<DegradeRule> rules = degradeRules.get(resource.getName());
    if (rules == null) {
        return;
    }

    for (DegradeRule rule : rules) {
        if (!rule.passCheck(context, node, count)) {
            throw new DegradeException(rule.getLimitApp(), rule);
        }
    }
}
 

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

@Override
public void entry(Context context, ResourceWrapper resourceWrapper, DefaultNode node, int count,
                  boolean prioritized, Object... args) throws Throwable {
    checkFlow(resourceWrapper, context, node, count, prioritized);

    fireEntry(context, resourceWrapper, node, count, prioritized, args);
}
 

private static boolean fallbackToLocalOrPass(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                             Object value) {
    if (rule.getClusterConfig().isFallbackToLocalWhenFail()) {
        return passLocalCheck(resourceWrapper, rule, count, value);
    } else {
        // The rule won't be activated, just pass.
        return true;
    }
}
 

@Override
public void exit(Context context, ResourceWrapper resourceWrapper, int count, Object... args) {
    try {
        fireExit(context, resourceWrapper, count, args);
    } catch (Throwable e) {
        RecordLog.warn("Unexpected entry exit exception", e);
    }
}
 

@Override
public void onExit(Context context, ResourceWrapper resourceWrapper, int count, Object... args) {
    for (MetricExtension m : MetricExtensionProvider.getMetricExtensions()) {
        if (context.getCurEntry().getError() == null) {
            long realRt = TimeUtil.currentTimeMillis() - context.getCurEntry().getCreateTime();
            m.addRt(resourceWrapper.getName(), realRt, args);
            m.addSuccess(resourceWrapper.getName(), count, args);
            m.decreaseThreadNum(resourceWrapper.getName(), args);
        }
    }
}