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

Example 11 with QuotaResponse

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

the class CalleeControllerTests method setUp.

@Before
public void setUp() {
    QuotaResponse quotaResponse = mock(QuotaResponse.class);
    when(quotaResponse.getCode()).thenReturn(QuotaResultCode.QuotaResultOk);
    when(limitAPI.getQuota(any())).thenReturn(quotaResponse);
}

Example 12 with QuotaResponse

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

the class QuotaCheckReactiveFilterTest method setUp.

@Before
public void setUp() {
    MetadataContext.LOCAL_NAMESPACE = "TEST";
    LimitAPI 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 {
            return new QuotaResponse(new QuotaResult(null, 0, null));
        }
    });
    PolarisRateLimitProperties polarisRateLimitProperties = new PolarisRateLimitProperties();
    polarisRateLimitProperties.setRejectRequestTips("RejectRequestTips");
    polarisRateLimitProperties.setRejectHttpCode(419);
    RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
    when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.EMPTY_SET);
    this.quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver);
}

Example 13 with QuotaResponse

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

the class QuotaCheckServletFilterTest method setUp.

@Before
public void setUp() {
    MetadataContext.LOCAL_NAMESPACE = "TEST";
    LimitAPI 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 {
            return new QuotaResponse(new QuotaResult(null, 0, null));
        }
    });
    PolarisRateLimitProperties polarisRateLimitProperties = new PolarisRateLimitProperties();
    polarisRateLimitProperties.setRejectRequestTips("RejectRequestTips");
    polarisRateLimitProperties.setRejectHttpCode(419);
    RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
    when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.EMPTY_SET);
    this.quotaCheckServletFilter = new QuotaCheckServletFilter(limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver);
}

Example 14 with QuotaResponse

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

the class QuotaCheckUtilsTest method testGetQuota.

@Test
public void testGetQuota() {
    // Pass
    String serviceName = "TestApp1";
    QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, null, null);
    assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
    assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
    assertThat(quotaResponse.getInfo()).isEqualTo("QuotaResultOk");
    // Unirate waiting 1000ms
    serviceName = "TestApp2";
    quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, null, null);
    assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
    assertThat(quotaResponse.getWaitMs()).isEqualTo(1000);
    assertThat(quotaResponse.getInfo()).isEqualTo("QuotaResultOk");
    // Rate limited
    serviceName = "TestApp3";
    quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, null, null);
    assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultLimited);
    assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
    assertThat(quotaResponse.getInfo()).isEqualTo("QuotaResultLimited");
    // Exception
    serviceName = "TestApp4";
    quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, null, null);
    assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
    assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
    assertThat(quotaResponse.getInfo()).isEqualTo("get quota failed");
}

Example 15 with QuotaResponse

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

the class QuotaCheckReactiveFilter method filter.

@Override
public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
    String localNamespace = MetadataContext.LOCAL_NAMESPACE;
    String localService = MetadataContext.LOCAL_SERVICE;
    Map<String, String> labels = getRequestLabels(exchange, localNamespace, localService);
    try {
        String path = exchange.getRequest().getURI().getPath();
        QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, localNamespace, localService, 1, labels, path);
        if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
            ServerHttpResponse response = exchange.getResponse();
            response.setRawStatusCode(polarisRateLimitProperties.getRejectHttpCode());
            response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
            DataBuffer dataBuffer = response.bufferFactory().allocateBuffer().write(rejectTips.getBytes(StandardCharsets.UTF_8));
            return response.writeWith(Mono.just(dataBuffer));
        }
        // Unirate
        if (quotaResponse.getCode() == QuotaResultCode.QuotaResultOk && quotaResponse.getWaitMs() > 0) {
            Thread.sleep(quotaResponse.getWaitMs());
        }
    } 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);
    }
    return chain.filter(exchange);
}