Examples with QuotaResponse com.tencent.polaris.ratelimit.api.rpc.QuotaResponse used on opensource projects

Example 6 with QuotaResponse

use of com.tencent.polaris.ratelimit.api.rpc.QuotaResponse in project spring-cloud-tencent by Tencent.

the class RateLimitZuulFilter method run.

@Override
public Object run() {
    // get request context
    RequestContext requestContext = RequestContext.getCurrentContext();
    String peerNamespace = MetadataContextHolder.get().getSystemMetadata(MetadataConstant.SystemMetadataKey.PEER_NAMESPACE);
    String peerService = MetadataContextHolder.get().getSystemMetadata(MetadataConstant.SystemMetadataKey.PEER_SERVICE);
    String peerPath = MetadataContextHolder.get().getSystemMetadata(MetadataConstant.SystemMetadataKey.PEER_PATH);
    Map<String, String> labels = null;
    if (StringUtils.isNotBlank(peerPath)) {
        labels = new HashMap<>();
        labels.put("method", peerPath);
    }
    try {
        QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, peerNamespace, peerService, 1, labels, null);
        if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
            requestContext.setSendZuulResponse(false);
            requestContext.setResponseStatusCode(TOO_MANY_REQUESTS.value());
            requestContext.getResponse().getWriter().write(Consts.QUOTA_LIMITED_INFO + quotaResponse.getInfo());
        }
    } catch (Throwable throwable) {
        // 限流API调用出现异常，不应该影响业务流程的调用
        LOG.error("fail to rate limit with QuotaRequest[{}-{}-{}].", peerNamespace, peerService, peerPath, throwable);
    }
    return null;
}

Example 7 with QuotaResponse

use of com.tencent.polaris.ratelimit.api.rpc.QuotaResponse in project spring-cloud-tencent by Tencent.

the class CalleeControllerTests method test1.

@Test
public void test1() {
    String url = "http://localhost:" + port + "/test/info";
    boolean hasPassed = false;
    boolean hasLimited = false;
    for (int i = 0; i < 30; i++) {
        try {
            if (i > 9) {
                QuotaResponse quotaResponse = mock(QuotaResponse.class);
                when(quotaResponse.getCode()).thenReturn(QuotaResultCode.QuotaResultLimited);
                when(quotaResponse.getInfo()).thenReturn("Testing rate limit after 10 times success.");
                when(limitAPI.getQuota(any())).thenReturn(quotaResponse);
            }
            String result = restTemplate.getForObject(url, String.class);
            System.out.println(result + " [" + i + "]");
            hasPassed = true;
        } catch (RestClientException e) {
            if (e instanceof TooManyRequests) {
                System.out.println(((TooManyRequests) e).getResponseBodyAsString());
                hasLimited = true;
            } else {
                e.printStackTrace();
                Assert.fail(e.getMessage());
            }
        }
    }
    Assert.assertTrue(hasPassed);
    Assert.assertTrue(hasLimited);
}

Example 8 with QuotaResponse

use of com.tencent.polaris.ratelimit.api.rpc.QuotaResponse in project spring-cloud-tencent by Tencent.

the class QuotaCheckUtilsTest method setUp.

@Before
public void setUp() {
    limitAPI = mock(LimitAPI.class);
    when(limitAPI.getQuota(any(QuotaRequest.class))).thenAnswer(invocationOnMock -> {
        String serviceName = ((QuotaRequest) invocationOnMock.getArgument(0)).getService();
        if (serviceName.equals("TestApp1")) {
            return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultOk, 0, "QuotaResultOk"));
        } else if (serviceName.equals("TestApp2")) {
            return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultOk, 1000, "QuotaResultOk"));
        } else if (serviceName.equals("TestApp3")) {
            return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultLimited, 0, "QuotaResultLimited"));
        } else {
            throw new RuntimeException("Mock exception.");
        }
    });
}

Example 9 with QuotaResponse

use of com.tencent.polaris.ratelimit.api.rpc.QuotaResponse in project spring-cloud-tencent by Tencent.

the class QuotaCheckServletFilter method doFilterInternal.

@Override
protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
    String localNamespace = MetadataContext.LOCAL_NAMESPACE;
    String localService = MetadataContext.LOCAL_SERVICE;
    Map<String, String> labels = getRequestLabels(request, localNamespace, localService);
    try {
        QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, localNamespace, localService, 1, labels, request.getRequestURI());
        if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
            response.setStatus(polarisRateLimitProperties.getRejectHttpCode());
            response.getWriter().write(rejectTips);
            return;
        }
        // Unirate
        if (quotaResponse.getCode() == QuotaResultCode.QuotaResultOk && quotaResponse.getWaitMs() > 0) {
            Thread.sleep(quotaResponse.getWaitMs());
        }
        filterChain.doFilter(request, response);
    } catch (Throwable t) {
        // An exception occurs in the rate limiting API call,
        // which should not affect the call of the business process.
        LOG.error("fail to invoke getQuota, service is " + localService, t);
        filterChain.doFilter(request, response);
    }
}

Example 10 with QuotaResponse

use of com.tencent.polaris.ratelimit.api.rpc.QuotaResponse in project spring-cloud-tencent by Tencent.

the class QuotaCheckUtils method getQuota.

public static QuotaResponse getQuota(LimitAPI limitAPI, String namespace, String service, int count, Map<String, String> labels, String method) {
    // build quota request
    QuotaRequest quotaRequest = new QuotaRequest();
    quotaRequest.setNamespace(namespace);
    quotaRequest.setService(service);
    quotaRequest.setCount(count);
    quotaRequest.setLabels(labels);
    quotaRequest.setMethod(method);
    try {
        return limitAPI.getQuota(quotaRequest);
    } catch (Throwable throwable) {
        LOG.error("fail to invoke getQuota of LimitAPI with QuotaRequest[{}].", quotaRequest, throwable);
        return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultOk, 0, "get quota failed"));
    }
}