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

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

@Test(expected = AuthorityException.class)
public void testCheckAuthorityNoExceptionItemsBlackFail() throws Exception {
    String origin = "appA";
    String resourceName = "testCheckAuthorityNoExceptionItemsBlackFail";
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    ContextUtil.enter("entrance", origin);
    try {
        AuthorityRule ruleA = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp(origin + ",appC")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_BLACK);

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

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.getType(), 1);
        } else {
            entry = ctSph.entry(resourceWrapper, 1);
        }
    } catch (BlockException ex) {
        fail("Unexpected blocked: " + ex.getClass().getCanonicalName());
    } finally {
        if (entry != null) {
            entry.exit();
        }
        ContextUtil.exit();
    }
}
 

@Test(expected = AuthorityException.class)
public void testCheckAuthorityNoExceptionItemsBlackFail() throws Exception {
    String origin = "appA";
    String resourceName = "testCheckAuthorityNoExceptionItemsBlackFail";
    ResourceWrapper resourceWrapper = new StringResourceWrapper(resourceName, EntryType.IN);
    ContextUtil.enter("entrance", origin);
    try {
        AuthorityRule ruleA = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp(origin + ",appC")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_BLACK);

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

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

/**
 * 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();
    }
}
 

@Test(expected = IllegalStateException.class)
public void testCleanCurrentEntryInLocalError() {
    final String contextName = "abc";
    try {
        ContextUtil.enter(contextName);
        Context curContext = ContextUtil.getContext();
        AsyncEntry entry = new AsyncEntry(new StringResourceWrapper("testCleanCurrentEntryInLocal", EntryType.OUT),
            null, curContext);

        entry.cleanCurrentEntryInLocal();

        entry.cleanCurrentEntryInLocal();
    } finally {
        ContextTestUtil.cleanUpContext();
    }
}
 

private static void doSomething(MessageExt message) {
    pool.submit(() -> {
        Entry entry = null;
        try {
            ContextUtil.enter(KEY);
            entry = SphU.entry(KEY, EntryType.OUT);

            // Your business logic here.
            System.out.printf("[%d][%s][Success: %d] Receive New Messages: %s %n", System.currentTimeMillis(),
                Thread.currentThread().getName(), SUCCESS_COUNT.addAndGet(1), new String(message.getBody()));
        } catch (BlockException ex) {
            // Blocked.
            System.out.println("Blocked: " + FAIL_COUNT.addAndGet(1));
        } finally {
            if (entry != null) {
                entry.exit();
            }
            ContextUtil.exit();
        }
    });
}
 

private void entryWhenSubscribed() {
    ContextConfig sentinelContextConfig = entryConfig.getContextConfig();
    if (sentinelContextConfig != null) {
        // If current we're already in a context, the context config won't work.
        ContextUtil.enter(sentinelContextConfig.getContextName(), sentinelContextConfig.getOrigin());
    }
    try {
        AsyncEntry entry = SphU.asyncEntry(entryConfig.getResourceName(), entryConfig.getEntryType(),
            entryConfig.getAcquireCount(), entryConfig.getArgs());
        this.currentEntry = entry;
        actual.onSubscribe(this);
    } catch (BlockException ex) {
        // Mark as completed (exited) explicitly.
        entryExited.set(true);
        // Signal cancel and propagate the {@code BlockException}.
        cancel();
        actual.onSubscribe(this);
        actual.onError(ex);
    } finally {
        if (sentinelContextConfig != null) {
            ContextUtil.exit();
        }
    }
}
 

@Test
public void testTraceWhenContextSizeExceedsThreshold() {
    int i = 0;
    for (; i < Constants.MAX_CONTEXT_NAME_SIZE; i++) {
        ContextUtil.enter("test-context-" + i);
        ContextUtil.exit();
    }

    try {
        ContextUtil.enter("test-context-" + i);
        throw new RuntimeException("test");
    } catch (Exception e) {
        Tracer.trace(e);
    } finally {
        ContextUtil.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);
}
 

@Test
public void testInvokeAsync() {

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

    when(invocation.getAttachment(ASYNC_KEY)).thenReturn(Boolean.TRUE.toString());
    final Result result = mock(Result.class);
    when(result.hasException()).thenReturn(false);
    when(invoker.invoke(invocation)).thenAnswer(invocationOnMock -> {
        verifyInvocationStructureForAsyncCall(invoker, invocation);
         return result;
    });
    consumerFilter.invoke(invoker, invocation);
    verify(invoker).invoke(invocation);

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

private static void doSomething(MessageExt message) {
    pool.submit(() -> {
        Entry entry = null;
        try {
            ContextUtil.enter(KEY);
            entry = SphU.entry(KEY, EntryType.OUT);

            // Your business logic here.
            System.out.printf("[%d][%s][Success: %d] Receive New Messages: %s %n", System.currentTimeMillis(),
                Thread.currentThread().getName(), SUCCESS_COUNT.addAndGet(1), new String(message.getBody()));
        } catch (BlockException ex) {
            // Blocked.
            System.out.println("Blocked: " + FAIL_COUNT.addAndGet(1));
        } finally {
            if (entry != null) {
                entry.exit();
            }
            ContextUtil.exit();
        }
    });
}
 

@AfterClass
public static void shutdown() {
    ctx.close();

    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 testExitTwoLastEntriesWithCustomContext() {
    String contextName = "context-rpc";
    ContextUtil.enter(contextName);
    Context context = ContextUtil.getContext();
    try {
        CtEntry entry1 = new CtEntry(new StringResourceWrapper("resA", EntryType.IN),
            null, context);
        entry1.exit();
        assertEquals(context, ContextUtil.getContext());
        CtEntry entry2 = new CtEntry(new StringResourceWrapper("resB", EntryType.IN),
            null, context);
        entry2.exit();
        assertEquals(context, ContextUtil.getContext());
    } finally {
        ContextUtil.exit();
        assertNull(ContextUtil.getContext());
    }
}
 

@Test
public void testExitLastEntryWithDefaultContext() {
    final Context defaultContext = getFakeDefaultContext();
    ContextUtil.runOnContext(defaultContext, new Runnable() {
        @Override
        public void run() {
            CtEntry entry = new CtEntry(new StringResourceWrapper("res", EntryType.IN),
                null, ContextUtil.getContext());
            assertSame(entry, defaultContext.getCurEntry());
            assertSame(defaultContext, ContextUtil.getContext());
            entry.exit();
            assertNull(defaultContext.getCurEntry());
            // Default context will be automatically exited.
            assertNull(ContextUtil.getContext());
        }
    });

}
 

@Test
public void testStringEntryCount() {
    if (SphO.entry("resourceName", 2)) {
        try {
            assertTrue(StringUtil.equalsIgnoreCase(
                ContextUtil.getContext().getCurEntry().getResourceWrapper().getName(), "resourceName"));
            assertSame(ContextUtil.getContext().getCurEntry().getResourceWrapper().getType(), EntryType.OUT);
        } finally {
            SphO.exit(2);
        }
    }
}
 

@Test
public void testStringEntryCount() throws BlockException {
    Entry e = SphU.entry("resourceName", 2);

    assertNotNull(e);
    assertEquals("resourceName", e.resourceWrapper.getName());
    assertEquals(e.resourceWrapper.getType(), EntryType.OUT);
    assertEquals(ContextUtil.getContext().getName(), Constants.CONTEXT_DEFAULT_NAME);

    e.exit(2);
}
 

@Test
public void testMethodEntryCount() throws NoSuchMethodException, SecurityException {
    Method method = SphOTest.class.getMethod("testMethodEntryCount");
    if (SphO.entry(method, 2)) {
        try {
            assertTrue(StringUtil.equalsIgnoreCase(
                ContextUtil.getContext().getCurEntry().getResourceWrapper().getName(),
                "com.alibaba.csp.sentinel.SphOTest:testMethodEntryCount()"));
            assertSame(ContextUtil.getContext().getCurEntry().getResourceWrapper().getType(), EntryType.OUT);
        } finally {
            SphO.exit(2);
        }
    }
}
 

@Test
public void testStringEntryNormal() {
    if (SphO.entry("resourceName")) {
        try {
            assertTrue(StringUtil.equalsIgnoreCase(
                ContextUtil.getContext().getCurEntry().getResourceWrapper().getName(), "resourceName"));
        } finally {
            SphO.exit();
        }
    }
}
 

@Test
public void testStringEntryType() {
    if (SphO.entry("resourceName", EntryType.IN)) {
        try {
            assertTrue(StringUtil.equalsIgnoreCase(
                ContextUtil.getContext().getCurEntry().getResourceWrapper().getName(), "resourceName"));
            assertSame(ContextUtil.getContext().getCurEntry().getResourceWrapper().getType(), EntryType.IN);
        } finally {
            SphO.exit();
        }
    }
}
 

@Test
public void testMethodEntryNormal() throws NoSuchMethodException, SecurityException {
    Method method = SphOTest.class.getMethod("testMethodEntryNormal");
    if (SphO.entry(method)) {
        try {
            assertTrue(StringUtil.equalsIgnoreCase(
                ContextUtil.getContext().getCurEntry().getResourceWrapper().getName(),
                "com.alibaba.csp.sentinel.SphOTest:testMethodEntryNormal()"));
        } finally {
            SphO.exit();
        }
    }
}
 

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

public void testMultipleLayer() throws Exception {
    // TODO: fix this
    ContextUtil.enter("entry1", "appA");

    Entry nodeA = SphU.entry("nodeA");
    assertSame(ContextUtil.getContext().getCurEntry(), nodeA);

    DefaultNode dnA = (DefaultNode)nodeA.getCurNode();
    assertNotNull(dnA);
    assertSame("nodeA", dnA.getId().getName());

    Entry nodeB = SphU.entry("nodeB");
    assertSame(ContextUtil.getContext().getCurEntry(), nodeB);
    DefaultNode dnB = (DefaultNode)nodeB.getCurNode();
    assertNotNull(dnB);
    assertTrue(dnA.getChildList().contains(dnB));

    Entry nodeC = SphU.entry("nodeC");
    assertSame(ContextUtil.getContext().getCurEntry(), nodeC);
    DefaultNode dnC = (DefaultNode)nodeC.getCurNode();
    assertNotNull(dnC);
    assertTrue(dnB.getChildList().contains(dnC));

    if (nodeC != null) {
        nodeC.exit();
    }
    assertSame(ContextUtil.getContext().getCurEntry(), nodeB);

    if (nodeB != null) {
        nodeB.exit();
    }
    assertSame(ContextUtil.getContext().getCurEntry(), nodeA);

    if (nodeA != null) {
        nodeA.exit();
    }
    assertNull(ContextUtil.getContext().getCurEntry());
    ContextUtil.exit();

}
 

@Test
public void testMethodEntryNormal() throws BlockException, NoSuchMethodException, SecurityException {
    Method method = SphUTest.class.getMethod("testMethodEntryNormal");
    Entry e = SphU.entry(method);

    assertNotNull(e);
    assertTrue(StringUtil
        .equalsIgnoreCase(e.resourceWrapper.getName(),
            "com.alibaba.csp.sentinel.SphUTest:testMethodEntryNormal()"));
    assertEquals(e.resourceWrapper.getType(), EntryType.OUT);
    assertEquals(ContextUtil.getContext().getName(), Constants.CONTEXT_DEFAULT_NAME);

    e.exit();
}
 

@Test
public void testStringEntryTypeCount() {
    if (SphO.entry("resourceName", EntryType.IN, 2)) {
        try {
            assertTrue(StringUtil.equalsIgnoreCase(
                ContextUtil.getContext().getCurEntry().getResourceWrapper().getName(), "resourceName"));
            assertSame(ContextUtil.getContext().getCurEntry().getResourceWrapper().getType(), EntryType.IN);
        } finally {
            SphO.exit(2);
        }
    }
}
 

@SuppressWarnings("unchecked")
static void tryExitFromCurrentContext() {
    RequestContext ctx = RequestContext.getCurrentContext();
    if (ctx.containsKey(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY)) {
        Deque<EntryHolder> holders = (Deque<EntryHolder>) ctx.get(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY);
        EntryHolder holder;
        while (!holders.isEmpty()) {
            holder = holders.pop();
            exit(holder);
        }
        ctx.remove(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY);
    }

    ContextUtil.exit();
}
 

@SuppressWarnings("unchecked")
static void tryTraceExceptionThenExitFromCurrentContext(Throwable t) {
    RequestContext ctx = RequestContext.getCurrentContext();
    if (ctx.containsKey(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY)) {
        Deque<EntryHolder> holders = (Deque<EntryHolder>) ctx.get(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY);
        EntryHolder holder;
        while (!holders.isEmpty()) {
            holder = holders.pop();
            Tracer.traceEntry(t, holder.getEntry());
            exit(holder);
        }
        ctx.remove(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY);
    }
    ContextUtil.exit();
}