com.alibaba.csp.sentinel.context.Context Java Examples

/**
 * Clean up resources.
 */
protected static void cleanUpAll() {
    Context context = ContextUtil.getContext();
    if (context != null) {
        context.setCurEntry(null);
        ContextUtil.exit();
    }

    Constants.ROOT.removeChildList();

    ClusterBuilderSlot.getClusterNodeMap().clear();

    // Clear chainMap in CtSph
    try {
        Method resetChainMapMethod = CtSph.class.getDeclaredMethod("resetChainMap");
        resetChainMapMethod.setAccessible(true);
        resetChainMapMethod.invoke(null);
    } catch (Exception e) {
        // Empty
    }
}
 

@Test
public void testCleanCurrentEntryInLocal() {
    final String contextName = "abc";
    try {
        ContextUtil.enter(contextName);
        Context curContext = ContextUtil.getContext();
        Entry previousEntry = new CtEntry(new StringResourceWrapper("entry-sync", EntryType.IN),
            null, curContext);
        AsyncEntry entry = new AsyncEntry(new StringResourceWrapper("testCleanCurrentEntryInLocal", EntryType.OUT),
            null, curContext);

        assertSame(entry, curContext.getCurEntry());

        entry.cleanCurrentEntryInLocal();
        assertNotSame(entry, curContext.getCurEntry());
        assertSame(previousEntry, curContext.getCurEntry());
    } finally {
        ContextTestUtil.cleanUpContext();
    }
}
 

void checkBlackWhiteAuthority(ResourceWrapper resource, Context context) throws AuthorityException {
    Map<String, Set<AuthorityRule>> authorityRules = AuthorityRuleManager.getAuthorityRules();

    if (authorityRules == null) {
        return;
    }

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

    for (AuthorityRule rule : rules) {
        if (!AuthorityRuleChecker.passCheck(rule, context)) {
            throw new AuthorityException(context.getOrigin(), rule);
        }
    }
}
 

@Test
public void testInvokeSync() {

    Invocation invocation = DubboTestUtil.getDefaultMockInvocationOne();
    Invoker invoker = DubboTestUtil.getDefaultMockInvoker();

    final Result result = mock(Result.class);
    when(result.hasException()).thenReturn(false);
    when(result.getException()).thenReturn(new Exception());
    when(invoker.invoke(invocation)).thenAnswer(invocationOnMock -> {
        verifyInvocationStructure(invoker, invocation);
        return result;
    });

    consumerFilter.invoke(invoker, invocation);
    verify(invoker).invoke(invocation);

    Context context = ContextUtil.getContext();
    assertNull(context);
}
 

@Test
public void testInvoke() {
    final Invoker invoker = mock(Invoker.class);
    when(invoker.getInterface()).thenReturn(DemoService.class);

    final Invocation invocation = mock(Invocation.class);
    Method method = DemoService.class.getMethods()[0];
    when(invocation.getMethodName()).thenReturn(method.getName());
    when(invocation.getParameterTypes()).thenReturn(method.getParameterTypes());

    final Result result = mock(Result.class);
    when(result.hasException()).thenReturn(false);
    when(invoker.invoke(invocation)).thenAnswer(new Answer<Object>() {
        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            verifyInvocationStructure(invoker, invocation);
            return result;
        }
    });

    filter.invoke(invoker, invocation);
    verify(invoker).invoke(invocation);

    Context context = ContextUtil.getContext();
    assertNull(context);
}
 

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

@Test
public void testSelectReferenceNodeForContextEntrance() {
    String contextName = "good_context";

    DefaultNode node = mock(DefaultNode.class);
    Context context = mock(Context.class);

    FlowRule rule = new FlowRule("testSelectReferenceNodeForContextEntrance")
        .setCount(1)
        .setStrategy(RuleConstant.STRATEGY_CHAIN)
        .setRefResource(contextName);

    when(context.getName()).thenReturn(contextName);
    assertEquals(node, FlowRuleChecker.selectReferenceNode(rule, context, node));

    when(context.getName()).thenReturn("other_context");
    assertNull(FlowRuleChecker.selectReferenceNode(rule, context, node));
}
 

private void testEntryAmountExceeded(boolean async) {
    fillFullResources();
    Entry entry = null;
    try {
        if (!async) {
            entry = ctSph.entry("testSync", EntryType.IN, 1);
        } else {
            entry = ctSph.asyncEntry("testSync", EntryType.IN, 1);
        }
        assertNull(((CtEntry)entry).chain);
        if (!async) {
            assertSame(entry, ContextUtil.getContext().getCurEntry());
        } else {
            Context asyncContext = ((AsyncEntry)entry).getAsyncContext();
            assertNotNull(asyncContext);
            assertSame(entry, asyncContext.getCurEntry());
        }
    } catch (BlockException ex) {
        fail("Unexpected blocked: " + ex.getClass().getCanonicalName());
    } finally {
        if (entry != null) {
            entry.exit();
        }
    }
}
 

@Test
public void testInvoke() {
    final Invoker invoker = mock(Invoker.class);
    when(invoker.getInterface()).thenReturn(DemoService.class);

    final Invocation invocation = mock(Invocation.class);
    Method method = DemoService.class.getMethods()[0];
    when(invocation.getMethodName()).thenReturn(method.getName());
    when(invocation.getParameterTypes()).thenReturn(method.getParameterTypes());

    final Result result = mock(Result.class);
    when(result.hasException()).thenReturn(false);
    when(invoker.invoke(invocation)).thenAnswer(invocationOnMock -> {
        verifyInvocationStructure(invoker, invocation);
        return result;
    });

    filter.invoke(invoker, invocation);
    verify(invoker).invoke(invocation);

    Context context = ContextUtil.getContext();
    assertNull(context);
}
 

static Node selectReferenceNode(FlowRule rule, Context context, DefaultNode node) {
    String refResource = rule.getRefResource();
    int strategy = rule.getStrategy();

    if (StringUtil.isEmpty(refResource)) {
        return null;
    }

    if (strategy == RuleConstant.STRATEGY_RELATE) {
        return ClusterBuilderSlot.getClusterNode(refResource);
    }

    if (strategy == RuleConstant.STRATEGY_CHAIN) {
        if (!refResource.equals(context.getName())) {
            return null;
        }
        return node;
    }
    // No node.
    return null;
}
 

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

        Throwable ex = context.getCurEntry().getError();
        if (ex != null) {
            m.addException(resource, count, ex);
        }
    }
}
 

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

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

static Node selectNodeByRequesterAndStrategy(/*@NonNull*/ FlowRule rule, Context context, DefaultNode node) {
    // The limit app should not be empty.
    String limitApp = rule.getLimitApp();
    int strategy = rule.getStrategy();
    String origin = context.getOrigin();

    if (limitApp.equals(origin) && filterOrigin(origin)) {
        if (strategy == RuleConstant.STRATEGY_DIRECT) {
            // Matches limit origin, return origin statistic node.
            return context.getOriginNode();
        }

        return selectReferenceNode(rule, context, node);
    } else if (RuleConstant.LIMIT_APP_DEFAULT.equals(limitApp)) {
        if (strategy == RuleConstant.STRATEGY_DIRECT) {
            // Return the cluster node.
            return node.getClusterNode();
        }

        return selectReferenceNode(rule, context, node);
    } else if (RuleConstant.LIMIT_APP_OTHER.equals(limitApp)
        && FlowRuleManager.isOtherOrigin(origin, rule.getResource())) {
        if (strategy == RuleConstant.STRATEGY_DIRECT) {
            return context.getOriginNode();
        }

        return selectReferenceNode(rule, context, node);
    }

    return null;
}
 

private static boolean fallbackToLocalOrPass(FlowRule rule, Context context, DefaultNode node, int acquireCount,
                                             boolean prioritized) {
    if (rule.getClusterConfig().isFallbackToLocalWhenFail()) {
        return passLocalCheck(rule, context, node, acquireCount, prioritized);
    } else {
        // The rule won't be activated, just pass.
        return true;
    }
}
 

@Test
public void testSelectNodeForRelateReference() {
    String refResource = "testSelectNodeForRelateReference_refResource";

    DefaultNode node = mock(DefaultNode.class);
    ClusterNode refCn = mock(ClusterNode.class);
    ClusterBuilderSlot.getClusterNodeMap().put(new StringResourceWrapper(refResource, EntryType.IN), refCn);
    Context context = mock(Context.class);

    FlowRule rule = new FlowRule("testSelectNodeForRelateReference")
        .setCount(1)
        .setStrategy(RuleConstant.STRATEGY_RELATE)
        .setRefResource(refResource);
    assertEquals(refCn, FlowRuleChecker.selectReferenceNode(rule, context, node));
}
 

@Test
public void testDuplicateInitAsyncContext() {
    Context context = new Context(null, "abc");
    AsyncEntry entry = new AsyncEntry(new StringResourceWrapper("testDuplicateInitAsyncContext", EntryType.OUT),
        null, context);
    entry.initAsyncContext();
    Context asyncContext = entry.getAsyncContext();

    // Duplicate init.
    entry.initAsyncContext();
    assertSame(asyncContext, entry.getAsyncContext());
}
 

@Test
public void testDefaultLimitAppFlowSelectNode() {
    DefaultNode node = mock(DefaultNode.class);
    ClusterNode cn = mock(ClusterNode.class);
    when(node.getClusterNode()).thenReturn(cn);
    Context context = mock(Context.class);

    // limitApp: default
    FlowRule rule = new FlowRule("testDefaultLimitAppFlowSelectNode").setCount(1);
    assertEquals(cn, FlowRuleChecker.selectNodeByRequesterAndStrategy(rule, context, node));
}
 

@Test
public void testExitNotMatchCurEntry() {
    String contextName = "context-rpc";
    ContextUtil.enter(contextName);
    Context context = ContextUtil.getContext();
    CtEntry entry1 = null;
    CtEntry entry2 = null;
    try {
        entry1 = new CtEntry(new StringResourceWrapper("res1", EntryType.IN),
            null, ContextUtil.getContext());
        assertSame(entry1, context.getCurEntry());
        entry2 = new CtEntry(new StringResourceWrapper("res2", EntryType.IN),
            null, ContextUtil.getContext());
        assertSame(entry2, context.getCurEntry());

        // Forget to exit for entry 2...
        // Directly exit for entry 1, then boom...
        entry1.exit();
    } catch (ErrorEntryFreeException ex) {
        assertNotNull(entry1);
        assertNotNull(entry2);
        assertNull(entry1.context);
        assertNull(entry2.context);
        assertNull(context.getCurEntry());
        return;
    } finally {
        ContextUtil.exit();
    }
    fail("Mismatch entry-exit should throw an ErrorEntryFreeException");
}
 

@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) {
    if (context.getCurEntry().getBlockError() == null) {
        ParameterMetric parameterMetric = ParameterMetricStorage.getParamMetric(resourceWrapper);

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

private AsyncEntry asyncEntryWithNoChain(ResourceWrapper resourceWrapper, Context context) {
    AsyncEntry entry = new AsyncEntry(resourceWrapper, null, context);
    entry.initAsyncContext();
    // The async entry will be removed from current context as soon as it has been created.
    entry.cleanCurrentEntryInLocal();
    return entry;
}
 

public static <R> Mono<R> entryWith(String resourceName, EntryType entryType, Mono<R> actual) {
    final AtomicReference<AsyncEntry> entryWrapper = new AtomicReference<>(null);
    return Mono.defer(() -> {
        try {
            AsyncEntry entry = SphU.asyncEntry(resourceName, entryType);
            entryWrapper.set(entry);
            return actual.subscriberContext(context -> {
                if (entry == null) {
                    return context;
                }
                Context sentinelContext = entry.getAsyncContext();
                if (sentinelContext == null) {
                    return context;
                }
                // TODO: check GC friendly?
                return context.put(SentinelReactorConstants.SENTINEL_CONTEXT_KEY, sentinelContext);
            }).doOnSuccessOrError((o, t) -> {
                if (entry != null && entryWrapper.compareAndSet(entry, null)) {
                    if (t != null) {
                        Tracer.traceContext(t, 1, entry.getAsyncContext());
                    }
                    entry.exit();
                }
            });
        } catch (BlockException ex) {
            return Mono.error(ex);
        }
    });
}
 

@Test
public void testDefaultLimitAppFlowSelectNode() {
    DefaultNode node = mock(DefaultNode.class);
    ClusterNode cn = mock(ClusterNode.class);
    when(node.getClusterNode()).thenReturn(cn);
    Context context = mock(Context.class);

    // limitApp: default
    FlowRule rule = new FlowRule("testDefaultLimitAppFlowSelectNode").setCount(1);
    assertEquals(cn, FlowRuleChecker.selectNodeByRequesterAndStrategy(rule, context, node));
}
 

@Test
public void testDuplicateInitAsyncContext() {
    Context context = new Context(null, "abc");
    AsyncEntry entry = new AsyncEntry(new StringResourceWrapper("testDuplicateInitAsyncContext", EntryType.OUT),
        null, context);
    entry.initAsyncContext();
    Context asyncContext = entry.getAsyncContext();

    // Duplicate init.
    entry.initAsyncContext();
    assertSame(asyncContext, entry.getAsyncContext());
}
 

@Override
public void onPass(Context context, ResourceWrapper resourceWrapper, DefaultNode param,
                   int count, Object... args) throws Exception {
    for (MetricExtension m : MetricExtensionProvider.getMetricExtensions()) {
        m.increaseThreadNum(resourceWrapper.getName(), args);
        m.addPass(resourceWrapper.getName(), count, args);
    }
}
 

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

CtEntry(ResourceWrapper resourceWrapper, ProcessorSlot<Object> chain, Context context) {
    super(resourceWrapper);
    this.chain = chain;
    this.context = context;

    setUpEntryFor(context);
}
 

private void setUpEntryFor(Context context) {
    // The entry should not be associated to NullContext.
    if (context instanceof NullContext) {
        return;
    }
    this.parent = context.getCurEntry();
    if (parent != null) {
        ((CtEntry)parent).child = this;
    }
    context.setCurEntry(this);
}
 

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