Java Code Examples for com.alibaba.csp.sentinel.context.ContextUtil#enter()

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

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

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

@Override
public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
    throws Exception {
    try {
        String resourceName = getResourceName(request);

        if (StringUtil.isNotEmpty(resourceName)) {
            // Parse the request origin using registered origin parser.
            String origin = parseOrigin(request);
            String contextName = getContextName(request);
            ContextUtil.enter(contextName, origin);
            Entry entry = SphU.entry(resourceName, ResourceTypeConstants.COMMON_WEB, EntryType.IN);

            setEntryInRequest(request, baseWebMvcConfig.getRequestAttributeName(), entry);
        }
        return true;
    } catch (BlockException e) {
        try {
            handleBlockException(request, response, e);
        } finally {
            ContextUtil.exit();
        }
        return false;
    }
}
 

@Test
public void testInitAndGetAsyncContext() {
    final String contextName = "abc";
    final String origin = "xxx";
    try {
        ContextUtil.enter(contextName, origin);
        Context curContext = ContextUtil.getContext();
        AsyncEntry entry = new AsyncEntry(new StringResourceWrapper("testInitAndGetAsyncContext", EntryType.OUT),
            null, curContext);
        assertNull(entry.getAsyncContext());

        entry.initAsyncContext();

        Context asyncContext = entry.getAsyncContext();
        assertNotNull(asyncContext);
        assertEquals(contextName, asyncContext.getName());
        assertEquals(origin, asyncContext.getOrigin());
        assertSame(curContext.getEntranceNode(), asyncContext.getEntranceNode());
        assertSame(entry, asyncContext.getCurEntry());
        assertTrue(asyncContext.isAsync());
    } finally {
        ContextTestUtil.cleanUpContext();
    }
}
 

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

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

@Override
public <ReqT, RespT> Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> serverCall, Metadata metadata,
                                                  ServerCallHandler<ReqT, RespT> serverCallHandler) {
    String resourceName = serverCall.getMethodDescriptor().getFullMethodName();
    // Remote address: serverCall.getAttributes().get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR);
    Entry entry = null;
    try {
        ContextUtil.enter(resourceName);
        entry = SphU.entry(resourceName, EntryType.IN);
        // Allow access, forward the call.
        return new ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT>(
            serverCallHandler.startCall(
                new ForwardingServerCall.SimpleForwardingServerCall<ReqT, RespT>(serverCall) {
                    @Override
                    public void close(Status status, Metadata trailers) {
                        super.close(status, trailers);
                        // Record the exception metrics.
                        if (!status.isOk()) {
                            recordException(status.asRuntimeException());
                        }
                    }
                }, metadata)) {};
    } catch (BlockException e) {
        serverCall.close(FLOW_CONTROL_BLOCK, new Metadata());
        return new ServerCall.Listener<ReqT>() {};
    } finally {
        if (entry != null) {
            entry.exit();
        }
        ContextUtil.exit();
    }
}
 

@Test
public void clusterNodeBuilder_normal() throws Exception {
    ContextUtil.enter("entry1", "caller1");

    Entry nodeA = SphU.entry("nodeA");

    Node curNode = nodeA.getCurNode();
    assertSame(curNode.getClass(), DefaultNode.class);
    DefaultNode dN = (DefaultNode)curNode;
    assertTrue(dN.getClusterNode().getOriginCountMap().containsKey("caller1"));
    assertSame(nodeA.getOriginNode(), dN.getClusterNode().getOrCreateOriginNode("caller1"));

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

    ContextUtil.enter("entry4", "caller2");

    nodeA = SphU.entry("nodeA");

    curNode = nodeA.getCurNode();
    assertSame(curNode.getClass(), DefaultNode.class);
    DefaultNode dN1 = (DefaultNode)curNode;
    assertTrue(dN1.getClusterNode().getOriginCountMap().containsKey("caller2"));
    assertNotSame(dN1, dN);

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

@Test
public void clusterNodeBuilder_normal() throws Exception {
    ContextUtil.enter("entry1", "caller1");

    Entry nodeA = SphU.entry("nodeA");

    Node curNode = nodeA.getCurNode();
    assertSame(curNode.getClass(), DefaultNode.class);
    DefaultNode dN = (DefaultNode)curNode;
    assertTrue(dN.getClusterNode().getOriginCountMap().containsKey("caller1"));
    assertSame(nodeA.getOriginNode(), dN.getClusterNode().getOrCreateOriginNode("caller1"));

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

    ContextUtil.enter("entry4", "caller2");

    nodeA = SphU.entry("nodeA");

    curNode = nodeA.getCurNode();
    assertSame(curNode.getClass(), DefaultNode.class);
    DefaultNode dN1 = (DefaultNode)curNode;
    assertTrue(dN1.getClusterNode().getOriginCountMap().containsKey("caller2"));
    assertNotSame(dN1, dN);

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

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

        AuthorityRuleManager.loadRules(Collections.singletonList(ruleA));
        authoritySlot.checkBlackWhiteAuthority(resourceWrapper, ContextUtil.getContext());

        AuthorityRule ruleB = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp("appD")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_BLACK);

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

private static void testFor(/*@NonNull*/ String resource, /*@NonNull*/ String origin) {
    ContextUtil.enter(resource, origin);
    Entry entry = null;
    try {
        entry = SphU.entry(resource);
        System.out.println(String.format("Passed for resource %s, origin is %s", resource, origin));
    } catch (BlockException ex) {
        System.err.println(String.format("Blocked for resource %s, origin is %s", resource, origin));
    } finally {
        if (entry != null) {
            entry.exit();
        }
        ContextUtil.exit();
    }
}
 

@Test
public void testPassCheck() {
    String origin = "appA";
    ContextUtil.enter("entrance", origin);
    try {
        String resourceName = "testPassCheck";
        AuthorityRule ruleA = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp(origin + ",appB")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_WHITE);
        AuthorityRule ruleB = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp("appB")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_WHITE);
        AuthorityRule ruleC = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp(origin)
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_BLACK);
        AuthorityRule ruleD = new AuthorityRule()
            .setResource(resourceName)
            .setLimitApp("appC")
            .as(AuthorityRule.class)
            .setStrategy(RuleConstant.AUTHORITY_BLACK);

        assertTrue(AuthorityRuleChecker.passCheck(ruleA, ContextUtil.getContext()));
        assertFalse(AuthorityRuleChecker.passCheck(ruleB, ContextUtil.getContext()));
        assertFalse(AuthorityRuleChecker.passCheck(ruleC, ContextUtil.getContext()));
        assertTrue(AuthorityRuleChecker.passCheck(ruleD, ContextUtil.getContext()));
    } finally {
        ContextUtil.exit();
    }
}
 

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

private Observable<HttpRequestMessage> apply(HttpRequestMessage request) {
    SessionContext context = request.getContext();
    Deque<EntryHolder> holders = new ArrayDeque<>();
    String routeId = routeExtractor.apply(request);
    String fallBackRoute = routeId;
    try {
        if (StringUtil.isNotBlank(routeId)) {
            ContextUtil.enter(GATEWAY_CONTEXT_ROUTE_PREFIX + routeId);
            doSentinelEntry(routeId, RESOURCE_MODE_ROUTE_ID, request, holders);
        }
        Set<String> matchingApis = pickMatchingApiDefinitions(request);
        if (!matchingApis.isEmpty() && ContextUtil.getContext() == null) {
            ContextUtil.enter(SentinelZuul2Constants.ZUUL_DEFAULT_CONTEXT);
        }
        for (String apiName : matchingApis) {
            fallBackRoute = apiName;
            doSentinelEntry(apiName, RESOURCE_MODE_CUSTOM_API_NAME, request, holders);
        }
        return Observable.just(request);
    } catch (BlockException t) {
        context.put(SentinelZuul2Constants.ZUUL_CTX_SENTINEL_BLOCKED_FLAG, Boolean.TRUE);
        context.put(SentinelZuul2Constants.ZUUL_CTX_SENTINEL_FALLBACK_ROUTE, fallBackRoute);
        if (fastError) {
            context.setShouldSendErrorResponse(true);
            context.setErrorEndpoint(blockedEndpointName);
        } else {
            context.setEndpoint(blockedEndpointName);
        }
        return Observable.error(t);
    } finally {
        if (!holders.isEmpty()) {
            context.put(SentinelZuul2Constants.ZUUL_CTX_SENTINEL_ENTRIES_KEY, holders);
        }
        // clear context to avoid another request use incorrect context
        ContextUtil.exit();
    }
}
 

@Override
public SofaResponse invoke(FilterInvoker invoker, SofaRequest request) throws SofaRpcException {
    // Now only support sync invoke.
    if (request.getInvokeType() != null && !RpcConstants.INVOKER_TYPE_SYNC.equals(request.getInvokeType())) {
        return invoker.invoke(request);
    }

    String callerApp = getApplicationName(request);
    String interfaceResourceName = getInterfaceResourceName(request);
    String methodResourceName = getMethodResourceName(request);

    Entry interfaceEntry = null;
    Entry methodEntry = null;
    try {
        ContextUtil.enter(methodResourceName, callerApp);

        interfaceEntry = SphU.entry(interfaceResourceName, ResourceTypeConstants.COMMON_RPC, EntryType.IN);
        methodEntry = SphU.entry(methodResourceName, ResourceTypeConstants.COMMON_RPC,
            EntryType.IN, getMethodArguments(request));

        SofaResponse response = invoker.invoke(request);

        traceResponseException(response, interfaceEntry, methodEntry);
        return response;
    } catch (BlockException e) {
        return SofaRpcFallbackRegistry.getProviderFallback().handle(invoker, request, e);
    } catch (Throwable t) {
        throw traceOtherException(t, interfaceEntry, methodEntry);
    } finally {
        if (methodEntry != null) {
            methodEntry.exit(1, getMethodArguments(request));
        }
        if (interfaceEntry != null) {
            interfaceEntry.exit();
        }
        ContextUtil.exit();
    }
}
 

private void fillFullContext() {
    for (int i = 0; i < Constants.MAX_CONTEXT_NAME_SIZE; i++) {
        ContextUtil.enter("test-context-" + i);
        ContextUtil.exit();
    }
}
 

@Override
public Object run() throws ZuulException {
    RequestContext ctx = RequestContext.getCurrentContext();
    String origin = parseOrigin(ctx.getRequest());
    String routeId = (String)ctx.get(ZuulConstant.PROXY_ID_KEY);

    Deque<AsyncEntry> asyncEntries = new ArrayDeque<>();
    String fallBackRoute = routeId;
    try {
        if (StringUtil.isNotBlank(routeId)) {
            ContextUtil.enter(GATEWAY_CONTEXT_ROUTE_PREFIX + routeId, origin);
            doSentinelEntry(routeId, RESOURCE_MODE_ROUTE_ID, ctx, asyncEntries);
        }

        Set<String> matchingApis = pickMatchingApiDefinitions(ctx);
        if (!matchingApis.isEmpty() && ContextUtil.getContext() == null) {
            ContextUtil.enter(ZuulConstant.ZUUL_DEFAULT_CONTEXT, origin);
        }
        for (String apiName : matchingApis) {
            fallBackRoute = apiName;
            doSentinelEntry(apiName, RESOURCE_MODE_CUSTOM_API_NAME, ctx, asyncEntries);
        }
    } catch (BlockException ex) {
        ZuulBlockFallbackProvider zuulBlockFallbackProvider = ZuulBlockFallbackManager.getFallbackProvider(
            fallBackRoute);
        BlockResponse blockResponse = zuulBlockFallbackProvider.fallbackResponse(fallBackRoute, ex);
        // Prevent routing from running
        ctx.setRouteHost(null);
        ctx.set(ZuulConstant.SERVICE_ID_KEY, null);

        // Set fallback response.
        ctx.setResponseBody(blockResponse.toString());
        ctx.setResponseStatusCode(blockResponse.getCode());
        // Set Response ContentType
        ctx.getResponse().setContentType("application/json; charset=utf-8");
    } finally {
        // We don't exit the entry here. We need to exit the entries in post filter to record Rt correctly.
        // So here the entries will be carried in the request context.
        if (!asyncEntries.isEmpty()) {
            ctx.put(ZuulConstant.ZUUL_CTX_SENTINEL_ENTRIES_KEY, asyncEntries);
        }
    }
    return null;
}
 

private void fillFullContext() {
    for (int i = 0; i < Constants.MAX_CONTEXT_NAME_SIZE; i++) {
        ContextUtil.enter("test-context-" + i);
        ContextUtil.exit();
    }
}