feature: add ratelimit provider info and refactor ratelimit use arguments (#902)

pull/919/head
Shanyou Yu (Sean Yu) 2 years ago committed by GitHub
parent 2d3af42bfe
commit 098d85d7f1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -3,6 +3,7 @@
- [fix: fix log feign response stream close bug.](https://github.com/Tencent/spring-cloud-tencent/pull/896)
- [fix:remove the secondary report.](https://github.com/Tencent/spring-cloud-tencent/pull/900)
- [feature: add ratelimit provider info and refactor ratelimit use arguments.](https://github.com/Tencent/spring-cloud-tencent/pull/902)
- [fix:optimize instance circuit beaker.](https://github.com/Tencent/spring-cloud-tencent/pull/908)
- [fix:optimize multi service registration and discovery.](https://github.com/Tencent/spring-cloud-tencent/pull/914)
- [feature: improve circuit breaker usage.](https://github.com/Tencent/spring-cloud-tencent/pull/916)

@ -28,6 +28,7 @@ import com.tencent.cloud.metadata.core.DecodeTransferMetadataServletFilter;
import com.tencent.cloud.metadata.core.EncodeTransferMedataFeignInterceptor;
import com.tencent.cloud.metadata.core.EncodeTransferMedataRestTemplateInterceptor;
import com.tencent.cloud.metadata.core.EncodeTransferMedataScgFilter;
import com.tencent.cloud.metadata.core.EncodeTransferMedataWebClientFilter;
import org.springframework.beans.factory.SmartInitializingSingleton;
import org.springframework.beans.factory.annotation.Autowired;
@ -39,6 +40,7 @@ import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import static javax.servlet.DispatcherType.ASYNC;
import static javax.servlet.DispatcherType.ERROR;
@ -143,4 +145,26 @@ public class MetadataTransferAutoConfiguration {
});
}
}
/**
* Create when WebClient.Builder exists.
*/
@Configuration(proxyBeanMethods = false)
@ConditionalOnClass(name = "org.springframework.web.reactive.function.client.WebClient")
protected static class MetadataTransferWebClientConfig {
@Autowired(required = false)
private List<WebClient.Builder> webClientBuilder = Collections.emptyList();
@Bean
public EncodeTransferMedataWebClientFilter encodeTransferMedataWebClientFilter() {
return new EncodeTransferMedataWebClientFilter();
}
@Bean
public SmartInitializingSingleton addEncodeTransferMetadataFilterForWebClient(EncodeTransferMedataWebClientFilter filter) {
return () -> webClientBuilder.forEach(webClient -> {
webClient.filter(filter);
});
}
}
}

@ -0,0 +1,90 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*
*/
package com.tencent.cloud.metadata.core;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.Map;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.common.util.JacksonUtils;
import reactor.core.publisher.Mono;
import org.springframework.util.CollectionUtils;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.ExchangeFunction;
import static com.tencent.cloud.common.constant.ContextConstant.UTF_8;
import static com.tencent.cloud.common.constant.MetadataConstant.HeaderName.CUSTOM_DISPOSABLE_METADATA;
import static com.tencent.cloud.common.constant.MetadataConstant.HeaderName.CUSTOM_METADATA;
/**
* web client filter used for writing metadata in HTTP request header.
*
* @author sean yu
*/
public class EncodeTransferMedataWebClientFilter implements ExchangeFilterFunction {
@Override
public Mono<ClientResponse> filter(ClientRequest clientRequest, ExchangeFunction next) {
MetadataContext metadataContext = MetadataContextHolder.get();
Map<String, String> customMetadata = metadataContext.getCustomMetadata();
Map<String, String> disposableMetadata = metadataContext.getDisposableMetadata();
Map<String, String> transHeaders = metadataContext.getTransHeadersKV();
ClientRequest.Builder requestBuilder = ClientRequest.from(clientRequest);
this.buildMetadataHeader(requestBuilder, customMetadata, CUSTOM_METADATA);
this.buildMetadataHeader(requestBuilder, disposableMetadata, CUSTOM_DISPOSABLE_METADATA);
this.buildTransmittedHeader(requestBuilder, transHeaders);
ClientRequest request = requestBuilder.build();
return next.exchange(request);
}
private void buildTransmittedHeader(ClientRequest.Builder requestBuilder, Map<String, String> transHeaders) {
if (!CollectionUtils.isEmpty(transHeaders)) {
transHeaders.forEach(requestBuilder::header);
}
}
/**
* Set metadata into the request header for {@link ClientRequest} .
* @param requestBuilder instance of {@link ClientRequest.Builder}
* @param metadata metadata map .
* @param headerName target metadata http header name .
*/
private void buildMetadataHeader(ClientRequest.Builder requestBuilder, Map<String, String> metadata, String headerName) {
if (!CollectionUtils.isEmpty(metadata)) {
String encodedMetadata = JacksonUtils.serialize2Json(metadata);
try {
requestBuilder.header(headerName, URLEncoder.encode(encodedMetadata, UTF_8));
}
catch (UnsupportedEncodingException e) {
requestBuilder.header(headerName, encodedMetadata);
}
}
}
}

@ -26,6 +26,7 @@ import java.util.stream.Collectors;
import com.tencent.cloud.metadata.core.EncodeTransferMedataFeignInterceptor;
import com.tencent.cloud.metadata.core.EncodeTransferMedataRestTemplateInterceptor;
import com.tencent.cloud.metadata.core.EncodeTransferMedataWebClientFilter;
import org.junit.jupiter.api.Test;
import org.springframework.boot.autoconfigure.AutoConfigurations;
@ -61,6 +62,7 @@ public class MetadataTransferAutoConfigurationTest {
assertThat(context).hasSingleBean(EncodeTransferMedataRestTemplateInterceptor.class);
assertThat(context).hasSingleBean(MetadataTransferAutoConfiguration.MetadataTransferScgFilterConfig.class);
assertThat(context).hasSingleBean(GlobalFilter.class);
assertThat(context).hasSingleBean(EncodeTransferMedataWebClientFilter.class);
});
}

@ -0,0 +1,84 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*
*/
package com.tencent.cloud.metadata.core;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.web.server.LocalServerPort;
import org.springframework.context.annotation.Bean;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT;
/**
* Test for {@link EncodeTransferMedataWebClientFilter}.
*
* @author sean yu
*/
@ExtendWith(SpringExtension.class)
@SpringBootTest(webEnvironment = RANDOM_PORT,
classes = EncodeTransferMedataWebClientFilterTest.TestApplication.class,
properties = {"spring.config.location = classpath:application-test.yml"})
public class EncodeTransferMedataWebClientFilterTest {
@Autowired
private WebClient.Builder webClientBuilder;
@Test
public void testTransitiveMetadataFromApplicationConfig() {
MetadataContext metadataContext = MetadataContextHolder.get();
metadataContext.setTransHeadersKV("xxx", "xxx");
String metadata = webClientBuilder.baseUrl("http://localhost:" + localServerPort).build()
.get()
.uri("/test")
.retrieve()
.bodyToMono(String.class)
.block();
assertThat(metadata).isEqualTo("2");
}
@LocalServerPort
private int localServerPort;
@SpringBootApplication
@RestController
protected static class TestApplication {
@Bean
public WebClient.Builder webClientBuilder() {
return WebClient.builder();
}
@RequestMapping("/test")
public String test() {
return MetadataContextHolder.get().getContext(MetadataContext.FRAGMENT_TRANSITIVE, "b");
}
}
}

@ -40,7 +40,7 @@ public class PolarisRouterAutoConfigurationTest {
PolarisLoadBalancerTest.class,
PolarisContextAutoConfiguration.class,
PolarisLoadBalancerAutoConfiguration.class))
.withPropertyValues("spring.cloud.loadbalancer.configurations=polaris");
.withPropertyValues("spring.cloud.loadbalancer.configurations=polaris", "spring.application.name=test");
private final ApplicationContextRunner noPolarisContextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(

@ -37,7 +37,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Haotian Zhang
*/
@ExtendWith(SpringExtension.class)
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT, classes = OkHttpUtilTest.TestApplication.class)
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT, classes = OkHttpUtilTest.TestApplication.class, properties = {"spring.application.name=test", "spring.cloud.polaris.discovery.register=false"})
public class OkHttpUtilTest {
@LocalServerPort

@ -19,6 +19,11 @@
<groupId>com.tencent.cloud</groupId>
<artifactId>spring-cloud-tencent-rpc-enhancement</artifactId>
</dependency>
<dependency>
<groupId>com.tencent.cloud</groupId>
<artifactId>spring-cloud-starter-tencent-metadata-transfer</artifactId>
</dependency>
<!-- Spring Cloud Tencent dependencies end -->
<!-- Polaris dependencies start -->

@ -36,6 +36,7 @@ import org.springframework.util.CollectionUtils;
*
*@author lepdou 2022-05-13
*/
@Deprecated
public class RateLimitRuleLabelResolver {
private final ServiceRuleManager serviceRuleManager;

@ -20,10 +20,10 @@ package com.tencent.cloud.polaris.ratelimit.config;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.context.config.PolarisContextAutoConfiguration;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckReactiveFilter;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilter;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
@ -31,6 +31,7 @@ import com.tencent.polaris.client.api.SDKContext;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.factory.LimitAPIFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
@ -39,6 +40,7 @@ import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.lang.Nullable;
import static com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant.FILTER_ORDER;
import static com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilter.QUOTA_FILTER_BEAN_NAME;
import static javax.servlet.DispatcherType.ASYNC;
import static javax.servlet.DispatcherType.ERROR;
@ -62,11 +64,6 @@ public class PolarisRateLimitAutoConfiguration {
return LimitAPIFactory.createLimitAPIByContext(polarisContext);
}
@Bean
public RateLimitRuleLabelResolver rateLimitRuleLabelService(ServiceRuleManager serviceRuleManager) {
return new RateLimitRuleLabelResolver(serviceRuleManager);
}
/**
* Create when web application type is SERVLET.
*/
@ -74,15 +71,19 @@ public class PolarisRateLimitAutoConfiguration {
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.SERVLET)
protected static class QuotaCheckFilterConfig {
@Bean
public RateLimitRuleArgumentServletResolver rateLimitRuleArgumentResolver(ServiceRuleManager serviceRuleManager,
@Autowired(required = false) PolarisRateLimiterLabelServletResolver labelResolver) {
return new RateLimitRuleArgumentServletResolver(serviceRuleManager, labelResolver);
}
@Bean
@ConditionalOnMissingBean
public QuotaCheckServletFilter quotaCheckFilter(LimitAPI limitAPI,
@Nullable PolarisRateLimiterLabelServletResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
return new QuotaCheckServletFilter(limitAPI, labelResolver,
polarisRateLimitProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
RateLimitRuleArgumentServletResolver rateLimitRuleArgumentResolver,
@Autowired(required = false) PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
return new QuotaCheckServletFilter(limitAPI, polarisRateLimitProperties, rateLimitRuleArgumentResolver, polarisRateLimiterLimitedFallback);
}
@Bean
@ -92,9 +93,11 @@ public class PolarisRateLimitAutoConfiguration {
quotaCheckServletFilter);
registrationBean.setDispatcherTypes(ASYNC, ERROR, FORWARD, INCLUDE, REQUEST);
registrationBean.setName(QUOTA_FILTER_BEAN_NAME);
registrationBean.setOrder(RateLimitConstant.FILTER_ORDER);
registrationBean.setOrder(FILTER_ORDER);
return registrationBean;
}
}
/**
@ -104,14 +107,18 @@ public class PolarisRateLimitAutoConfiguration {
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.REACTIVE)
protected static class MetadataReactiveFilterConfig {
@Bean
public RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentResolver(ServiceRuleManager serviceRuleManager,
@Autowired(required = false) PolarisRateLimiterLabelReactiveResolver labelResolver) {
return new RateLimitRuleArgumentReactiveResolver(serviceRuleManager, labelResolver);
}
@Bean
public QuotaCheckReactiveFilter quotaCheckReactiveFilter(LimitAPI limitAPI,
@Nullable PolarisRateLimiterLabelReactiveResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
return new QuotaCheckReactiveFilter(limitAPI, labelResolver,
polarisRateLimitProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
return new QuotaCheckReactiveFilter(limitAPI, polarisRateLimitProperties, rateLimitRuleArgumentResolver, polarisRateLimiterLimitedFallback);
}
}
}

@ -20,27 +20,22 @@ package com.tencent.cloud.polaris.ratelimit.filter;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import javax.annotation.PostConstruct;
import com.google.common.collect.Maps;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.util.expresstion.SpringWebExpressionLabelUtils;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils;
import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResultCode;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
@ -54,8 +49,6 @@ import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import static com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant.LABEL_METHOD;
/**
* Reactive filter to check quota.
*
@ -67,11 +60,9 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
private final LimitAPI limitAPI;
private final PolarisRateLimiterLabelReactiveResolver labelResolver;
private final PolarisRateLimitProperties polarisRateLimitProperties;
private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver;
private final RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentResolver;
private final PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
@ -79,14 +70,12 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
private String rejectTips;
public QuotaCheckReactiveFilter(LimitAPI limitAPI,
PolarisRateLimiterLabelReactiveResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
this.limitAPI = limitAPI;
this.labelResolver = labelResolver;
this.polarisRateLimitProperties = polarisRateLimitProperties;
this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver;
this.rateLimitRuleArgumentResolver = rateLimitRuleArgumentResolver;
this.polarisRateLimiterLimitedFallback = polarisRateLimiterLimitedFallback;
}
@ -105,12 +94,12 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
String localNamespace = MetadataContext.LOCAL_NAMESPACE;
String localService = MetadataContext.LOCAL_SERVICE;
Map<String, String> labels = getRequestLabels(exchange, localNamespace, localService);
Set<Argument> arguments = rateLimitRuleArgumentResolver.getArguments(exchange, localNamespace, localService);
long waitMs = -1;
try {
String path = exchange.getRequest().getURI().getPath();
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI,
localNamespace, localService, 1, labels, path);
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(
limitAPI, localNamespace, localService, 1, arguments, path);
if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
ServerHttpResponse response = exchange.getResponse();
@ -149,40 +138,4 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
}
}
private Map<String, String> getRequestLabels(ServerWebExchange exchange, String localNamespace, String localService) {
Map<String, String> labels = new HashMap<>();
// add build in labels
String path = exchange.getRequest().getURI().getPath();
if (StringUtils.isNotBlank(path)) {
labels.put(LABEL_METHOD, path);
}
// add rule expression labels
Map<String, String> expressionLabels = getRuleExpressionLabels(exchange, localNamespace, localService);
labels.putAll(expressionLabels);
// add custom labels
Map<String, String> customResolvedLabels = getCustomResolvedLabels(exchange);
labels.putAll(customResolvedLabels);
return labels;
}
private Map<String, String> getCustomResolvedLabels(ServerWebExchange exchange) {
if (labelResolver != null) {
try {
return labelResolver.resolve(exchange);
}
catch (Throwable e) {
LOGGER.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e);
}
}
return Maps.newHashMap();
}
private Map<String, String> getRuleExpressionLabels(ServerWebExchange exchange, String namespace, String service) {
Set<String> expressionLabels = rateLimitRuleLabelResolver.getExpressionLabelKeys(namespace, service);
return SpringWebExpressionLabelUtils.resolve(exchange, expressionLabels);
}
}

@ -19,9 +19,6 @@
package com.tencent.cloud.polaris.ratelimit.filter;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
@ -32,18 +29,16 @@ import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.util.expresstion.ServletExpressionLabelUtils;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils;
import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResultCode;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@ -53,8 +48,6 @@ import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.web.filter.OncePerRequestFilter;
import static com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant.LABEL_METHOD;
/**
* Servlet filter to check quota.
*
@ -70,25 +63,21 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
private static final Logger LOG = LoggerFactory.getLogger(QuotaCheckServletFilter.class);
private final LimitAPI limitAPI;
private final PolarisRateLimiterLabelServletResolver labelResolver;
private final PolarisRateLimitProperties polarisRateLimitProperties;
private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver;
private final RateLimitRuleArgumentServletResolver rateLimitRuleArgumentResolver;
private final PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
private String rejectTips;
public QuotaCheckServletFilter(LimitAPI limitAPI,
PolarisRateLimiterLabelServletResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
RateLimitRuleArgumentServletResolver rateLimitRuleArgumentResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
this.limitAPI = limitAPI;
this.labelResolver = labelResolver;
this.polarisRateLimitProperties = polarisRateLimitProperties;
this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver;
this.rateLimitRuleArgumentResolver = rateLimitRuleArgumentResolver;
this.polarisRateLimiterLimitedFallback = polarisRateLimiterLimitedFallback;
}
@ -104,11 +93,11 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
String localNamespace = MetadataContext.LOCAL_NAMESPACE;
String localService = MetadataContext.LOCAL_SERVICE;
Map<String, String> labels = getRequestLabels(request, localNamespace, localService);
Set<Argument> arguments = rateLimitRuleArgumentResolver.getArguments(request, localNamespace, localService);
try {
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI,
localNamespace, localService, 1, labels, request.getRequestURI());
localNamespace, localService, 1, arguments, request.getRequestURI());
if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
if (!Objects.isNull(polarisRateLimiterLimitedFallback)) {
@ -140,40 +129,4 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
filterChain.doFilter(request, response);
}
private Map<String, String> getRequestLabels(HttpServletRequest request, String localNamespace, String localService) {
Map<String, String> labels = new HashMap<>();
// add build in labels
String path = request.getRequestURI();
if (StringUtils.isNotBlank(path)) {
labels.put(LABEL_METHOD, path);
}
// add rule expression labels
Map<String, String> expressionLabels = getRuleExpressionLabels(request, localNamespace, localService);
labels.putAll(expressionLabels);
// add custom resolved labels
Map<String, String> customLabels = getCustomResolvedLabels(request);
labels.putAll(customLabels);
return labels;
}
private Map<String, String> getCustomResolvedLabels(HttpServletRequest request) {
if (labelResolver != null) {
try {
return labelResolver.resolve(request);
}
catch (Throwable e) {
LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e);
}
}
return Collections.emptyMap();
}
private Map<String, String> getRuleExpressionLabels(HttpServletRequest request, String namespace, String service) {
Set<String> expressionLabels = rateLimitRuleLabelResolver.getExpressionLabelKeys(namespace, service);
return ServletExpressionLabelUtils.resolve(request, expressionLabels);
}
}

@ -0,0 +1,123 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*
*/
package com.tencent.cloud.polaris.ratelimit.resolver;
import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.web.server.ServerWebExchange;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
/**
* resolve arguments from rate limit rule for Reactive.
*
* @author seansyyu 2023-03-09
*/
public class RateLimitRuleArgumentReactiveResolver {
private static final Logger LOG = LoggerFactory.getLogger(RateLimitRuleArgumentReactiveResolver.class);
private final ServiceRuleManager serviceRuleManager;
private final PolarisRateLimiterLabelReactiveResolver labelResolver;
public RateLimitRuleArgumentReactiveResolver(ServiceRuleManager serviceRuleManager, PolarisRateLimiterLabelReactiveResolver labelResolver) {
this.serviceRuleManager = serviceRuleManager;
this.labelResolver = labelResolver;
}
public Set<Argument> getArguments(ServerWebExchange request, String namespace, String service) {
RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service);
if (rateLimitRule == null) {
return Collections.emptySet();
}
List<RateLimitProto.Rule> rules = rateLimitRule.getRulesList();
if (CollectionUtils.isEmpty(rules)) {
return Collections.emptySet();
}
return rules.stream()
.flatMap(rule -> rule.getArgumentsList().stream())
.map(matchArgument -> {
String matchKey = matchArgument.getKey();
Argument argument = null;
switch (matchArgument.getType()) {
case CUSTOM:
argument = StringUtils.isBlank(matchKey) ? null :
Argument.buildCustom(matchKey, Optional.ofNullable(getCustomResolvedLabels(request).get(matchKey)).orElse(StringUtils.EMPTY));
break;
case METHOD:
argument = Argument.buildMethod(request.getRequest().getMethodValue());
break;
case HEADER:
argument = StringUtils.isBlank(matchKey) ? null :
Argument.buildHeader(matchKey, Optional.ofNullable(request.getRequest().getHeaders().getFirst(matchKey)).orElse(StringUtils.EMPTY));
break;
case QUERY:
argument = StringUtils.isBlank(matchKey) ? null :
Argument.buildQuery(matchKey, Optional.ofNullable(request.getRequest().getQueryParams().getFirst(matchKey)).orElse(StringUtils.EMPTY));
break;
case CALLER_SERVICE:
String sourceServiceNamespace = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, true).orElse(StringUtils.EMPTY);
String sourceServiceName = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAME, true).orElse(StringUtils.EMPTY);
if (!StringUtils.isEmpty(sourceServiceNamespace) && !StringUtils.isEmpty(sourceServiceName)) {
argument = Argument.buildCallerService(sourceServiceNamespace, sourceServiceName);
}
break;
case CALLER_IP:
InetSocketAddress remoteAddress = request.getRequest().getRemoteAddress();
argument = Argument.buildCallerIP(remoteAddress != null ? remoteAddress.getAddress().getHostAddress() : StringUtils.EMPTY);
break;
default:
break;
}
return argument;
}).filter(Objects::nonNull).collect(Collectors.toSet());
}
private Map<String, String> getCustomResolvedLabels(ServerWebExchange request) {
if (labelResolver != null) {
try {
return labelResolver.resolve(request);
}
catch (Throwable e) {
LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e);
}
}
return Collections.emptyMap();
}
}

@ -0,0 +1,123 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*
*/
package com.tencent.cloud.polaris.ratelimit.resolver;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import javax.servlet.http.HttpServletRequest;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilter;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
/**
* resolve arguments from rate limit rule for Servlet.
*
* @author seansyyu 2023-03-09
*/
public class RateLimitRuleArgumentServletResolver {
private static final Logger LOG = LoggerFactory.getLogger(QuotaCheckServletFilter.class);
private final ServiceRuleManager serviceRuleManager;
private final PolarisRateLimiterLabelServletResolver labelResolver;
public RateLimitRuleArgumentServletResolver(ServiceRuleManager serviceRuleManager, PolarisRateLimiterLabelServletResolver labelResolver) {
this.serviceRuleManager = serviceRuleManager;
this.labelResolver = labelResolver;
}
public Set<Argument> getArguments(HttpServletRequest request, String namespace, String service) {
RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service);
if (rateLimitRule == null) {
return Collections.emptySet();
}
List<RateLimitProto.Rule> rules = rateLimitRule.getRulesList();
if (CollectionUtils.isEmpty(rules)) {
return Collections.emptySet();
}
return rules.stream()
.flatMap(rule -> rule.getArgumentsList().stream())
.map(matchArgument -> {
String matchKey = matchArgument.getKey();
Argument argument = null;
switch (matchArgument.getType()) {
case CUSTOM:
argument = StringUtils.isBlank(matchKey) ? null :
Argument.buildCustom(matchKey, Optional.ofNullable(getCustomResolvedLabels(request).get(matchKey)).orElse(StringUtils.EMPTY));
break;
case METHOD:
argument = Argument.buildMethod(request.getMethod());
break;
case HEADER:
argument = StringUtils.isBlank(matchKey) ? null :
Argument.buildHeader(matchKey, Optional.ofNullable(request.getHeader(matchKey)).orElse(StringUtils.EMPTY));
break;
case QUERY:
argument = StringUtils.isBlank(matchKey) ? null :
Argument.buildQuery(matchKey, Optional.ofNullable(request.getParameter(matchKey)).orElse(StringUtils.EMPTY));
break;
case CALLER_SERVICE:
String sourceServiceNamespace = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, true).orElse(StringUtils.EMPTY);
String sourceServiceName = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAME, true).orElse(StringUtils.EMPTY);
if (!StringUtils.isEmpty(sourceServiceNamespace) && !StringUtils.isEmpty(sourceServiceName)) {
argument = Argument.buildCallerService(sourceServiceNamespace, sourceServiceName);
}
break;
case CALLER_IP:
argument = Argument.buildCallerIP(Optional.ofNullable(request.getRemoteAddr()).orElse(StringUtils.EMPTY));
break;
default:
break;
}
return argument;
}).filter(Objects::nonNull).collect(Collectors.toSet());
}
private Map<String, String> getCustomResolvedLabels(HttpServletRequest request) {
if (labelResolver != null) {
try {
return labelResolver.resolve(request);
}
catch (Throwable e) {
LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e);
}
}
return Collections.emptyMap();
}
}

@ -18,9 +18,11 @@
package com.tencent.cloud.polaris.ratelimit.utils;
import java.util.Map;
import java.util.Set;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse;
import org.slf4j.Logger;
@ -38,6 +40,7 @@ public final class QuotaCheckUtils {
private QuotaCheckUtils() {
}
@Deprecated
public static QuotaResponse getQuota(LimitAPI limitAPI, String namespace, String service, int count,
Map<String, String> labels, String method) {
// build quota request
@ -56,4 +59,23 @@ public final class QuotaCheckUtils {
return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultOk, 0, "get quota failed"));
}
}
public static QuotaResponse getQuota(LimitAPI limitAPI, String namespace, String service, int count,
Set<Argument> arguments, String method) {
// build quota request
QuotaRequest quotaRequest = new QuotaRequest();
quotaRequest.setNamespace(namespace);
quotaRequest.setService(service);
quotaRequest.setCount(count);
quotaRequest.setArguments(arguments);
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"));
}
}
}

@ -18,9 +18,10 @@
package com.tencent.cloud.polaris.ratelimit.config;
import com.tencent.cloud.polaris.context.config.PolarisContextAutoConfiguration;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckReactiveFilter;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilter;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentServletResolver;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import org.junit.jupiter.api.Test;
@ -55,7 +56,8 @@ public class PolarisRateLimitAutoConfigurationTest {
PolarisRateLimitAutoConfiguration.class))
.run(context -> {
assertThat(context).hasSingleBean(LimitAPI.class);
assertThat(context).hasSingleBean(RateLimitRuleLabelResolver.class);
assertThat(context).doesNotHaveBean(RateLimitRuleArgumentServletResolver.class);
assertThat(context).doesNotHaveBean(RateLimitRuleArgumentReactiveResolver.class);
assertThat(context).doesNotHaveBean(PolarisRateLimitAutoConfiguration.QuotaCheckFilterConfig.class);
assertThat(context).doesNotHaveBean(QuotaCheckServletFilter.class);
assertThat(context).doesNotHaveBean(FilterRegistrationBean.class);
@ -72,12 +74,13 @@ public class PolarisRateLimitAutoConfigurationTest {
PolarisRateLimitAutoConfiguration.class))
.run(context -> {
assertThat(context).hasSingleBean(LimitAPI.class);
assertThat(context).hasSingleBean(RateLimitRuleLabelResolver.class);
assertThat(context).hasSingleBean(RateLimitRuleArgumentServletResolver.class);
assertThat(context).hasSingleBean(PolarisRateLimitAutoConfiguration.QuotaCheckFilterConfig.class);
assertThat(context).hasSingleBean(QuotaCheckServletFilter.class);
assertThat(context).hasSingleBean(FilterRegistrationBean.class);
assertThat(context).doesNotHaveBean(PolarisRateLimitAutoConfiguration.MetadataReactiveFilterConfig.class);
assertThat(context).doesNotHaveBean(QuotaCheckReactiveFilter.class);
assertThat(context).doesNotHaveBean(RateLimitRuleArgumentReactiveResolver.class);
});
}
@ -89,7 +92,8 @@ public class PolarisRateLimitAutoConfigurationTest {
PolarisRateLimitAutoConfiguration.class))
.run(context -> {
assertThat(context).hasSingleBean(LimitAPI.class);
assertThat(context).hasSingleBean(RateLimitRuleLabelResolver.class);
assertThat(context).doesNotHaveBean(RateLimitRuleArgumentServletResolver.class);
assertThat(context).hasSingleBean(RateLimitRuleArgumentReactiveResolver.class);
assertThat(context).doesNotHaveBean(PolarisRateLimitAutoConfiguration.QuotaCheckFilterConfig.class);
assertThat(context).doesNotHaveBean(QuotaCheckServletFilter.class);
assertThat(context).doesNotHaveBean(FilterRegistrationBean.class);

@ -34,7 +34,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Haotian Zhang
*/
@ExtendWith(SpringExtension.class)
@SpringBootTest(classes = PolarisRateLimitPropertiesTest.TestApplication.class)
@SpringBootTest(classes = PolarisRateLimitPropertiesTest.TestApplication.class, properties = "spring.application.name=test")
@ActiveProfiles("test")
public class PolarisRateLimitPropertiesTest {

@ -17,32 +17,32 @@
package com.tencent.cloud.polaris.ratelimit.filter;
import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import com.tencent.cloud.common.util.expresstion.SpringWebExpressionLabelUtils;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
@ -60,10 +60,8 @@ import org.springframework.web.server.WebFilterChain;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anySet;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;
/**
@ -76,36 +74,15 @@ import static org.mockito.Mockito.when;
@SpringBootTest(classes = QuotaCheckReactiveFilterTest.TestApplication.class,
properties = {"spring.cloud.polaris.namespace=Test", "spring.cloud.polaris.service=TestApp"})
public class QuotaCheckReactiveFilterTest {
private static MockedStatic<ApplicationContextAwareUtils> mockedApplicationContextAwareUtils;
private static MockedStatic<SpringWebExpressionLabelUtils> expressionLabelUtilsMockedStatic;
private final PolarisRateLimiterLabelReactiveResolver labelResolver =
exchange -> Collections.singletonMap("ReactiveResolver", "ReactiveResolver");
exchange -> Collections.singletonMap("xxx", "xxx");
private QuotaCheckReactiveFilter quotaCheckReactiveFilter;
private QuotaCheckReactiveFilter quotaCheckWithRateLimiterLimitedFallbackReactiveFilter;
private PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
@BeforeAll
static void beforeAll() {
expressionLabelUtilsMockedStatic = mockStatic(SpringWebExpressionLabelUtils.class);
when(SpringWebExpressionLabelUtils.resolve(any(ServerWebExchange.class), anySet()))
.thenReturn(Collections.singletonMap("RuleLabelResolver", "RuleLabelResolver"));
mockedApplicationContextAwareUtils = Mockito.mockStatic(ApplicationContextAwareUtils.class);
mockedApplicationContextAwareUtils.when(() -> ApplicationContextAwareUtils.getProperties(anyString()))
.thenReturn("unit-test");
}
@AfterAll
static void afterAll() {
mockedApplicationContextAwareUtils.close();
expressionLabelUtilsMockedStatic.close();
}
@BeforeEach
void setUp() {
void setUp() throws InvalidProtocolBufferException {
MetadataContext.LOCAL_NAMESPACE = "TEST";
polarisRateLimiterLimitedFallback = new JsonPolarisRateLimiterLimitedFallback();
LimitAPI limitAPI = mock(LimitAPI.class);
when(limitAPI.getQuota(any(QuotaRequest.class))).thenAnswer(invocationOnMock -> {
@ -128,14 +105,24 @@ public class QuotaCheckReactiveFilterTest {
polarisRateLimitProperties.setRejectRequestTips("RejectRequestTips提示消息");
polarisRateLimitProperties.setRejectHttpCode(419);
RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString()))
.thenReturn(Collections.emptySet());
PolarisRateLimitProperties polarisRateLimitWithHtmlRejectTipsProperties = new PolarisRateLimitProperties();
polarisRateLimitWithHtmlRejectTipsProperties.setRejectRequestTips("<h1>RejectRequestTips提示消息</h1>");
polarisRateLimitWithHtmlRejectTipsProperties.setRejectHttpCode(419);
this.quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(
limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver, null);
this.quotaCheckWithRateLimiterLimitedFallbackReactiveFilter = new QuotaCheckReactiveFilter(
limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
RateLimitProto.Rule.Builder ratelimitRuleBuilder = RateLimitProto.Rule.newBuilder();
InputStream inputStream = QuotaCheckServletFilterTest.class.getClassLoader().getResourceAsStream("ratelimit.json");
String json = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)).lines().collect(Collectors.joining(""));
JsonFormat.parser().ignoringUnknownFields().merge(json, ratelimitRuleBuilder);
RateLimitProto.Rule rateLimitRule = ratelimitRuleBuilder.build();
RateLimitProto.RateLimit rateLimit = RateLimitProto.RateLimit.newBuilder().addRules(rateLimitRule).build();
when(serviceRuleManager.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit);
RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager, labelResolver);
this.quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(limitAPI, polarisRateLimitProperties, rateLimitRuleArgumentReactiveResolver, null);
this.polarisRateLimiterLimitedFallback = new JsonPolarisRateLimiterLimitedFallback();
this.quotaCheckWithRateLimiterLimitedFallbackReactiveFilter = new QuotaCheckReactiveFilter(limitAPI, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleArgumentReactiveResolver, polarisRateLimiterLimitedFallback);
}
@Test
@ -156,42 +143,6 @@ public class QuotaCheckReactiveFilterTest {
}
}
@Test
public void testGetRuleExpressionLabels() {
try {
Method getCustomResolvedLabels =
QuotaCheckReactiveFilter.class.getDeclaredMethod("getCustomResolvedLabels", ServerWebExchange.class);
getCustomResolvedLabels.setAccessible(true);
// Mock request
MockServerHttpRequest request = MockServerHttpRequest.get("http://localhost:8080/test").build();
ServerWebExchange exchange = MockServerWebExchange.from(request);
// labelResolver != null
Map<String, String> result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange);
assertThat(result.size()).isEqualTo(1);
assertThat(result.get("ReactiveResolver")).isEqualTo("ReactiveResolver");
// throw exception
PolarisRateLimiterLabelReactiveResolver exceptionLabelResolver = exchange1 -> {
throw new RuntimeException("Mock exception.");
};
quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, exceptionLabelResolver, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange);
assertThat(result.size()).isEqualTo(0);
// labelResolver == null
quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, null, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange);
assertThat(result.size()).isEqualTo(0);
getCustomResolvedLabels.setAccessible(false);
}
catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
fail("Exception encountered.", e);
}
}
@Test
public void testFilter() {
// Create mock WebFilterChain
@ -199,19 +150,20 @@ public class QuotaCheckReactiveFilterTest {
// Mock request
MockServerHttpRequest request = MockServerHttpRequest.get("http://localhost:8080/test").build();
ServerWebExchange exchange = MockServerWebExchange.from(request);
quotaCheckReactiveFilter.init();
// Pass
MetadataContext.LOCAL_SERVICE = "TestApp1";
quotaCheckReactiveFilter.filter(exchange, webFilterChain);
ServerWebExchange testApp1Exchange = MockServerWebExchange.from(request);
quotaCheckReactiveFilter.filter(testApp1Exchange, webFilterChain);
// Unirate waiting 1000ms
MetadataContext.LOCAL_SERVICE = "TestApp2";
ServerWebExchange testApp2Exchange = MockServerWebExchange.from(request);
long startTimestamp = System.currentTimeMillis();
CountDownLatch countDownLatch = new CountDownLatch(1);
quotaCheckReactiveFilter.filter(exchange, webFilterChain).subscribe(e -> {
quotaCheckReactiveFilter.filter(testApp2Exchange, webFilterChain).subscribe(e -> {
}, t -> {
}, countDownLatch::countDown);
try {
@ -224,14 +176,16 @@ public class QuotaCheckReactiveFilterTest {
// Rate limited
MetadataContext.LOCAL_SERVICE = "TestApp3";
quotaCheckReactiveFilter.filter(exchange, webFilterChain);
ServerHttpResponse response = exchange.getResponse();
ServerWebExchange testApp3Exchange = MockServerWebExchange.from(request);
quotaCheckReactiveFilter.filter(testApp3Exchange, webFilterChain);
ServerHttpResponse response = testApp3Exchange.getResponse();
assertThat(response.getRawStatusCode()).isEqualTo(419);
assertThat(response.getStatusCode()).isEqualTo(HttpStatus.INSUFFICIENT_SPACE_ON_RESOURCE);
// Exception
MetadataContext.LOCAL_SERVICE = "TestApp4";
quotaCheckReactiveFilter.filter(exchange, webFilterChain);
ServerWebExchange testApp4Exchange = MockServerWebExchange.from(request);
quotaCheckReactiveFilter.filter(testApp4Exchange, webFilterChain);
}
@Test

@ -17,53 +17,47 @@
package com.tencent.cloud.polaris.ratelimit.filter;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.Map;
import java.util.stream.Collectors;
import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import com.tencent.cloud.common.util.expresstion.SpringWebExpressionLabelUtils;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.resolver.RateLimitRuleArgumentServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest;
import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.http.MediaType;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.mock.web.MockHttpServletResponse;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anySet;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;
/**
@ -71,41 +65,23 @@ import static org.mockito.Mockito.when;
*
* @author Haotian Zhang, cheese8
*/
@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.LENIENT)
@SpringBootTest(classes = QuotaCheckServletFilterTest.TestApplication.class, properties = {
@ExtendWith(SpringExtension.class)
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
classes = QuotaCheckServletFilterTest.TestApplication.class,
properties = {
"spring.cloud.polaris.namespace=Test", "spring.cloud.polaris.service=TestApp"
})
public class QuotaCheckServletFilterTest {
private static MockedStatic<ApplicationContextAwareUtils> mockedApplicationContextAwareUtils;
private static MockedStatic<SpringWebExpressionLabelUtils> expressionLabelUtilsMockedStatic;
private final PolarisRateLimiterLabelServletResolver labelResolver =
exchange -> Collections.singletonMap("ServletResolver", "ServletResolver");
exchange -> Collections.singletonMap("xxx", "xxx");
private QuotaCheckServletFilter quotaCheckServletFilter;
private QuotaCheckServletFilter quotaCheckWithHtmlRejectTipsServletFilter;
private QuotaCheckServletFilter quotaCheckWithRateLimiterLimitedFallbackFilter;
private PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
@BeforeAll
static void beforeAll() {
expressionLabelUtilsMockedStatic = mockStatic(SpringWebExpressionLabelUtils.class);
when(SpringWebExpressionLabelUtils.resolve(any(ServerWebExchange.class), anySet()))
.thenReturn(Collections.singletonMap("RuleLabelResolver", "RuleLabelResolver"));
mockedApplicationContextAwareUtils = Mockito.mockStatic(ApplicationContextAwareUtils.class);
mockedApplicationContextAwareUtils.when(() -> ApplicationContextAwareUtils.getProperties(anyString()))
.thenReturn("unit-test");
}
@AfterAll
static void afterAll() {
mockedApplicationContextAwareUtils.close();
expressionLabelUtilsMockedStatic.close();
}
@BeforeEach
void setUp() {
void setUp() throws InvalidProtocolBufferException {
MetadataContext.LOCAL_NAMESPACE = "TEST";
LimitAPI limitAPI = mock(LimitAPI.class);
@ -133,15 +109,21 @@ public class QuotaCheckServletFilterTest {
polarisRateLimitWithHtmlRejectTipsProperties.setRejectRequestTips("<h1>RejectRequestTips提示消息</h1>");
polarisRateLimitWithHtmlRejectTipsProperties.setRejectHttpCode(419);
RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.emptySet());
this.quotaCheckServletFilter = new QuotaCheckServletFilter(limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver, null);
this.quotaCheckWithHtmlRejectTipsServletFilter = new QuotaCheckServletFilter(
limitAPI, labelResolver, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleLabelResolver, null);
polarisRateLimiterLimitedFallback = new JsonPolarisRateLimiterLimitedFallback();
this.quotaCheckWithRateLimiterLimitedFallbackFilter = new QuotaCheckServletFilter(
limitAPI, labelResolver, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
RateLimitProto.Rule.Builder ratelimitRuleBuilder = RateLimitProto.Rule.newBuilder();
InputStream inputStream = QuotaCheckServletFilterTest.class.getClassLoader().getResourceAsStream("ratelimit.json");
String json = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)).lines().collect(Collectors.joining(""));
JsonFormat.parser().ignoringUnknownFields().merge(json, ratelimitRuleBuilder);
RateLimitProto.Rule rateLimitRule = ratelimitRuleBuilder.build();
RateLimitProto.RateLimit rateLimit = RateLimitProto.RateLimit.newBuilder().addRules(rateLimitRule).build();
when(serviceRuleManager.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit);
RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver = new RateLimitRuleArgumentServletResolver(serviceRuleManager, labelResolver);
this.quotaCheckServletFilter = new QuotaCheckServletFilter(limitAPI, polarisRateLimitProperties, rateLimitRuleArgumentServletResolver, null);
this.quotaCheckWithHtmlRejectTipsServletFilter = new QuotaCheckServletFilter(limitAPI, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleArgumentServletResolver, null);
this.polarisRateLimiterLimitedFallback = new JsonPolarisRateLimiterLimitedFallback();
this.quotaCheckWithRateLimiterLimitedFallbackFilter = new QuotaCheckServletFilter(limitAPI, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleArgumentServletResolver, polarisRateLimiterLimitedFallback);
}
@Test
@ -167,40 +149,6 @@ public class QuotaCheckServletFilterTest {
quotaCheckWithRateLimiterLimitedFallbackFilter.init();
}
@Test
public void testGetRuleExpressionLabels() {
try {
Method getCustomResolvedLabels = QuotaCheckServletFilter.class.getDeclaredMethod("getCustomResolvedLabels", HttpServletRequest.class);
getCustomResolvedLabels.setAccessible(true);
// Mock request
MockHttpServletRequest request = new MockHttpServletRequest();
// labelResolver != null
Map<String, String> result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request);
assertThat(result.size()).isEqualTo(1);
assertThat(result.get("ServletResolver")).isEqualTo("ServletResolver");
// throw exception
PolarisRateLimiterLabelServletResolver exceptionLabelResolver = request1 -> {
throw new RuntimeException("Mock exception.");
};
quotaCheckServletFilter = new QuotaCheckServletFilter(null, exceptionLabelResolver, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request);
assertThat(result.size()).isEqualTo(0);
// labelResolver == null
quotaCheckServletFilter = new QuotaCheckServletFilter(null, null, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request);
assertThat(result.size()).isEqualTo(0);
getCustomResolvedLabels.setAccessible(false);
}
catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
fail("Exception encountered.", e);
}
}
@Test
public void testDoFilterInternal() {
// Create mock FilterChain
@ -210,34 +158,39 @@ public class QuotaCheckServletFilterTest {
// Mock request
MockHttpServletRequest request = new MockHttpServletRequest();
MockHttpServletResponse response = new MockHttpServletResponse();
quotaCheckServletFilter.init();
quotaCheckWithHtmlRejectTipsServletFilter.init();
try {
// Pass
MetadataContext.LOCAL_SERVICE = "TestApp1";
quotaCheckServletFilter.doFilterInternal(request, response, filterChain);
MockHttpServletResponse testApp1Response = new MockHttpServletResponse();
quotaCheckServletFilter.doFilterInternal(request, testApp1Response, filterChain);
// Unirate waiting 1000ms
MetadataContext.LOCAL_SERVICE = "TestApp2";
MockHttpServletResponse testApp2Response = new MockHttpServletResponse();
long startTimestamp = System.currentTimeMillis();
quotaCheckServletFilter.doFilterInternal(request, response, filterChain);
quotaCheckServletFilter.doFilterInternal(request, testApp2Response, filterChain);
assertThat(System.currentTimeMillis() - startTimestamp).isGreaterThanOrEqualTo(1000L);
// Rate limited
MetadataContext.LOCAL_SERVICE = "TestApp3";
quotaCheckServletFilter.doFilterInternal(request, response, filterChain);
assertThat(response.getStatus()).isEqualTo(419);
assertThat(response.getContentAsString()).isEqualTo("RejectRequestTips提示消息");
MockHttpServletResponse testApp3Response = new MockHttpServletResponse();
quotaCheckServletFilter.doFilterInternal(request, testApp3Response, filterChain);
assertThat(testApp3Response.getStatus()).isEqualTo(419);
assertThat(testApp3Response.getContentAsString()).isEqualTo("RejectRequestTips提示消息");
quotaCheckWithHtmlRejectTipsServletFilter.doFilterInternal(request, response, filterChain);
assertThat(response.getStatus()).isEqualTo(419);
assertThat(response.getContentAsString()).isEqualTo("RejectRequestTips提示消息");
MockHttpServletResponse testApp3Response2 = new MockHttpServletResponse();
quotaCheckWithHtmlRejectTipsServletFilter.doFilterInternal(request, testApp3Response2, filterChain);
assertThat(testApp3Response2.getStatus()).isEqualTo(419);
assertThat(testApp3Response2.getContentAsString()).isEqualTo("<h1>RejectRequestTips提示消息</h1>");
// Exception
MockHttpServletResponse testApp4Response = new MockHttpServletResponse();
MetadataContext.LOCAL_SERVICE = "TestApp4";
quotaCheckServletFilter.doFilterInternal(request, response, filterChain);
quotaCheckServletFilter.doFilterInternal(request, testApp4Response, filterChain);
}
catch (ServletException | IOException e) {
fail("Exception encountered.", e);
@ -252,31 +205,34 @@ public class QuotaCheckServletFilterTest {
// Mock request
MockHttpServletRequest request = new MockHttpServletRequest();
MockHttpServletResponse response = new MockHttpServletResponse();
quotaCheckWithRateLimiterLimitedFallbackFilter.init();
try {
// Pass
MetadataContext.LOCAL_SERVICE = "TestApp1";
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
MockHttpServletResponse testApp1response = new MockHttpServletResponse();
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, testApp1response, filterChain);
// Unirate waiting 1000ms
MetadataContext.LOCAL_SERVICE = "TestApp2";
MockHttpServletResponse testApp2response = new MockHttpServletResponse();
long startTimestamp = System.currentTimeMillis();
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, testApp2response, filterChain);
assertThat(System.currentTimeMillis() - startTimestamp).isGreaterThanOrEqualTo(1000L);
// Rate limited
MetadataContext.LOCAL_SERVICE = "TestApp3";
MockHttpServletResponse testApp3response = new MockHttpServletResponse();
String contentType = new MediaType(polarisRateLimiterLimitedFallback.mediaType(), polarisRateLimiterLimitedFallback.charset()).toString();
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
assertThat(response.getStatus()).isEqualTo(polarisRateLimiterLimitedFallback.rejectHttpCode());
assertThat(response.getContentAsString()).isEqualTo(polarisRateLimiterLimitedFallback.rejectTips());
assertThat(response.getContentType()).isEqualTo(contentType);
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, testApp3response, filterChain);
assertThat(testApp3response.getStatus()).isEqualTo(polarisRateLimiterLimitedFallback.rejectHttpCode());
assertThat(testApp3response.getContentAsString()).isEqualTo(polarisRateLimiterLimitedFallback.rejectTips());
assertThat(testApp3response.getContentType()).isEqualTo(contentType);
// Exception
MetadataContext.LOCAL_SERVICE = "TestApp4";
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
MockHttpServletResponse testApp4response = new MockHttpServletResponse();
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, testApp4response, filterChain);
}
catch (ServletException | IOException e) {
fail("Exception encountered.", e);

@ -0,0 +1,144 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/
package com.tencent.cloud.polaris.ratelimit.resolver;
import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilterTest;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
import org.springframework.mock.web.server.MockServerWebExchange;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import org.springframework.web.server.ServerWebExchange;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@ExtendWith(SpringExtension.class)
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
classes = RateLimitRuleArgumentReactiveResolverTest.TestApplication.class,
properties = {
"spring.cloud.polaris.namespace=Test", "spring.cloud.polaris.service=TestApp"
})
public class RateLimitRuleArgumentReactiveResolverTest {
private final PolarisRateLimiterLabelReactiveResolver labelResolver =
exchange -> Collections.singletonMap("xxx", "xxx");
private final PolarisRateLimiterLabelReactiveResolver labelResolverEx =
exchange -> {
throw new RuntimeException();
};
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver1;
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver2;
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver3;
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver4;
@BeforeEach
void setUp() throws InvalidProtocolBufferException {
MetadataContext.LOCAL_NAMESPACE = "TEST";
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
RateLimitProto.Rule.Builder ratelimitRuleBuilder = RateLimitProto.Rule.newBuilder();
InputStream inputStream = QuotaCheckServletFilterTest.class.getClassLoader().getResourceAsStream("ratelimit.json");
String json = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)).lines().collect(Collectors.joining(""));
JsonFormat.parser().ignoringUnknownFields().merge(json, ratelimitRuleBuilder);
RateLimitProto.Rule rateLimitRule = ratelimitRuleBuilder.build();
RateLimitProto.RateLimit rateLimit = RateLimitProto.RateLimit.newBuilder().addRules(rateLimitRule).build();
when(serviceRuleManager.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit);
// normal
this.rateLimitRuleArgumentReactiveResolver1 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager, labelResolver);
// ex
this.rateLimitRuleArgumentReactiveResolver2 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager, labelResolverEx);
// null
ServiceRuleManager serviceRuleManager1 = mock(ServiceRuleManager.class);
when(serviceRuleManager1.getServiceRateLimitRule(anyString(), anyString())).thenReturn(null);
this.rateLimitRuleArgumentReactiveResolver3 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager1, labelResolver);
// null 2
ServiceRuleManager serviceRuleManager2 = mock(ServiceRuleManager.class);
RateLimitProto.RateLimit rateLimit2 = RateLimitProto.RateLimit.newBuilder().build();
when(serviceRuleManager2.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit2);
this.rateLimitRuleArgumentReactiveResolver4 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager2, labelResolver);
}
@Test
public void testGetRuleArguments() {
// Mock request
MetadataContext.LOCAL_SERVICE = "Test";
// Mock request
MockServerHttpRequest request = MockServerHttpRequest.get("http://127.0.0.1:8080/test")
.remoteAddress(new InetSocketAddress("127.0.0.1", 8080))
.header("xxx", "xxx")
.queryParam("yyy", "yyy")
.build();
ServerWebExchange exchange = MockServerWebExchange.from(request);
MetadataContext metadataContext = new MetadataContext();
metadataContext.setUpstreamDisposableMetadata(new HashMap<String, String>() {{
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE);
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE);
}});
MetadataContextHolder.set(metadataContext);
Set<Argument> arguments = rateLimitRuleArgumentReactiveResolver1.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
Set<Argument> exceptRes = new HashSet<>();
exceptRes.add(Argument.buildMethod("GET"));
exceptRes.add(Argument.buildHeader("xxx", "xxx"));
exceptRes.add(Argument.buildQuery("yyy", "yyy"));
exceptRes.add(Argument.buildCallerIP("127.0.0.1"));
exceptRes.add(Argument.buildCustom("xxx", "xxx"));
exceptRes.add(Argument.buildCallerService(MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE));
assertThat(arguments).isEqualTo(exceptRes);
rateLimitRuleArgumentReactiveResolver2.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
rateLimitRuleArgumentReactiveResolver3.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
rateLimitRuleArgumentReactiveResolver4.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
}
@SpringBootApplication
protected static class TestApplication {
}
}

@ -0,0 +1,135 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/
package com.tencent.cloud.polaris.ratelimit.resolver;
import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilterTest;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.polaris.ratelimit.api.rpc.Argument;
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@ExtendWith(SpringExtension.class)
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
classes = RateLimitRuleArgumentServletResolverTest.TestApplication.class,
properties = {
"spring.cloud.polaris.namespace=Test", "spring.cloud.polaris.service=TestApp"
})
public class RateLimitRuleArgumentServletResolverTest {
private final PolarisRateLimiterLabelServletResolver labelResolver =
exchange -> Collections.singletonMap("xxx", "xxx");
private final PolarisRateLimiterLabelServletResolver labelResolverEx =
exchange -> {
throw new RuntimeException();
};
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver1;
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver2;
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver3;
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver4;
@BeforeEach
void setUp() throws InvalidProtocolBufferException {
MetadataContext.LOCAL_NAMESPACE = "TEST";
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
RateLimitProto.Rule.Builder ratelimitRuleBuilder = RateLimitProto.Rule.newBuilder();
InputStream inputStream = QuotaCheckServletFilterTest.class.getClassLoader().getResourceAsStream("ratelimit.json");
String json = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)).lines().collect(Collectors.joining(""));
JsonFormat.parser().ignoringUnknownFields().merge(json, ratelimitRuleBuilder);
RateLimitProto.Rule rateLimitRule = ratelimitRuleBuilder.build();
RateLimitProto.RateLimit rateLimit = RateLimitProto.RateLimit.newBuilder().addRules(rateLimitRule).build();
when(serviceRuleManager.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit);
// normal
this.rateLimitRuleArgumentServletResolver1 = new RateLimitRuleArgumentServletResolver(serviceRuleManager, labelResolver);
// ex
this.rateLimitRuleArgumentServletResolver2 = new RateLimitRuleArgumentServletResolver(serviceRuleManager, labelResolverEx);
// null
ServiceRuleManager serviceRuleManager1 = mock(ServiceRuleManager.class);
when(serviceRuleManager1.getServiceRateLimitRule(anyString(), anyString())).thenReturn(null);
this.rateLimitRuleArgumentServletResolver3 = new RateLimitRuleArgumentServletResolver(serviceRuleManager1, labelResolver);
// null 2
ServiceRuleManager serviceRuleManager2 = mock(ServiceRuleManager.class);
RateLimitProto.RateLimit rateLimit2 = RateLimitProto.RateLimit.newBuilder().build();
when(serviceRuleManager2.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit2);
this.rateLimitRuleArgumentServletResolver4 = new RateLimitRuleArgumentServletResolver(serviceRuleManager2, labelResolver);
}
@Test
public void testGetRuleArguments() {
// Mock request
MetadataContext.LOCAL_SERVICE = "Test";
MockHttpServletRequest request = new MockHttpServletRequest(null, "GET", "/xxx");
request.setParameter("yyy", "yyy");
request.addHeader("xxx", "xxx");
MetadataContext metadataContext = new MetadataContext();
metadataContext.setUpstreamDisposableMetadata(new HashMap<String, String>() {{
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE);
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE);
}});
MetadataContextHolder.set(metadataContext);
Set<Argument> arguments = rateLimitRuleArgumentServletResolver1.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
Set<Argument> exceptRes = new HashSet<>();
exceptRes.add(Argument.buildMethod("GET"));
exceptRes.add(Argument.buildHeader("xxx", "xxx"));
exceptRes.add(Argument.buildQuery("yyy", "yyy"));
exceptRes.add(Argument.buildCallerIP("127.0.0.1"));
exceptRes.add(Argument.buildCustom("xxx", "xxx"));
exceptRes.add(Argument.buildCallerService(MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE));
assertThat(arguments).isEqualTo(exceptRes);
rateLimitRuleArgumentServletResolver2.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
rateLimitRuleArgumentServletResolver3.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
rateLimitRuleArgumentServletResolver4.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
}
@SpringBootApplication
protected static class TestApplication {
}
}

@ -17,6 +17,9 @@
package com.tencent.cloud.polaris.ratelimit.utils;
import java.util.HashMap;
import java.util.HashSet;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest;
@ -66,28 +69,59 @@ public class QuotaCheckUtilsTest {
public void testGetQuota() {
// Pass
String serviceName = "TestApp1";
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, null, null);
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashMap<>(), 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, new HashMap<>(), 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, new HashMap<>(), 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, new HashMap<>(), null);
assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
assertThat(quotaResponse.getInfo()).isEqualTo("get quota failed");
}
@Test
public void testGetQuota2() {
// Pass
String serviceName = "TestApp1";
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashSet<>(), 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);
quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashSet<>(), 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);
quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashSet<>(), 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);
quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashSet<>(), null);
assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
assertThat(quotaResponse.getInfo()).isEqualTo("get quota failed");

@ -0,0 +1,95 @@
{
"id": "f7560cce829e4d4c8556a6be63539af5",
"service": "xxx",
"namespace": "default",
"subset": {},
"priority": 0,
"resource": "QPS",
"type": "GLOBAL",
"labels": {},
"amounts": [
{
"maxAmount": 1,
"validDuration": "1s",
"precision": null,
"startAmount": null,
"minAmount": null
}
],
"action": "REJECT",
"disable": false,
"report": null,
"ctime": "2022-12-11 21:56:59",
"mtime": "2023-03-10 15:40:33",
"revision": "6eec6f416bee40ecbf664c93add61358",
"service_token": null,
"adjuster": null,
"regex_combine": true,
"amountMode": "GLOBAL_TOTAL",
"failover": "FAILOVER_LOCAL",
"cluster": null,
"method": {
"type": "EXACT",
"value": "/xxx",
"value_type": "TEXT"
},
"arguments": [
{
"type": "CALLER_SERVICE",
"key": "default",
"value": {
"type": "EXACT",
"value": "xxx",
"value_type": "TEXT"
}
},
{
"type": "HEADER",
"key": "xxx",
"value": {
"type": "EXACT",
"value": "xxx",
"value_type": "TEXT"
}
},
{
"type": "QUERY",
"key": "yyy",
"value": {
"type": "EXACT",
"value": "yyy",
"value_type": "TEXT"
}
},
{
"type": "METHOD",
"key": "$method",
"value": {
"type": "EXACT",
"value": "GET",
"value_type": "TEXT"
}
},
{
"type": "CALLER_IP",
"key": "$caller_ip",
"value": {
"type": "EXACT",
"value": "127.0.0.1",
"value_type": "TEXT"
}
},
{
"type": "CUSTOM",
"key": "xxx",
"value": {
"type": "EXACT",
"value": "xxx",
"value_type": "TEXT"
}
}
],
"name": "xxx",
"etime": "2023-03-10 15:40:33",
"max_queue_delay": 30
}

@ -19,6 +19,7 @@
package com.tencent.cloud.polaris.router.config;
import com.tencent.cloud.common.metadata.config.MetadataAutoConfiguration;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import com.tencent.cloud.polaris.context.config.PolarisContextAutoConfiguration;
import com.tencent.cloud.polaris.router.feign.RouterLabelFeignInterceptor;
import org.junit.jupiter.api.Test;
@ -39,7 +40,9 @@ public class FeignAutoConfigurationTest {
MetadataAutoConfiguration.class,
RouterAutoConfiguration.class,
PolarisContextAutoConfiguration.class,
FeignAutoConfiguration.class));
FeignAutoConfiguration.class,
ApplicationContextAwareUtils.class
)).withPropertyValues("spring.application.name=test");
@Test
public void routerLabelInterceptor() {

@ -19,6 +19,7 @@
package com.tencent.cloud.polaris.router.config;
import com.tencent.cloud.common.metadata.config.MetadataAutoConfiguration;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import com.tencent.cloud.polaris.context.config.PolarisContextAutoConfiguration;
import org.junit.jupiter.api.Test;
@ -42,7 +43,9 @@ public class RouterAutoConfigurationTests {
MetadataAutoConfiguration.class,
RouterAutoConfiguration.class,
PolarisContextAutoConfiguration.class,
RouterAutoConfiguration.RouterLabelRestTemplateConfig.class));
RouterAutoConfiguration.RouterLabelRestTemplateConfig.class,
ApplicationContextAwareUtils.class
)).withPropertyValues("spring.application.name=test");
@Test
public void testRouterLabelRestTemplateConfig() {

@ -55,7 +55,7 @@ public final class MetadataConstant {
/**
* Order of filter.
*/
public static final int WEB_FILTER_ORDER = Ordered.HIGHEST_PRECEDENCE + 13;
public static final int WEB_FILTER_ORDER = Ordered.HIGHEST_PRECEDENCE + 9;
/**
* Order of MetadataFirstFeignInterceptor.
@ -93,4 +93,19 @@ public final class MetadataConstant {
*/
public static final String METADATA_CONTEXT = "SCT-METADATA-CONTEXT";
}
public static class DefaultMetadata {
/**
* Default Metadata Source Service Namespace Key.
*/
public static final String DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE = "source_service_namespace";
/**
* Default Metadata Source Service Name Key.
*/
public static final String DEFAULT_METADATA_SOURCE_SERVICE_NAME = "source_service_name";
}
}

@ -135,13 +135,14 @@ public final class MetadataContextHolder {
mergedTransitiveMetadata.putAll(staticTransitiveMetadata);
mergedTransitiveMetadata.putAll(dynamicTransitiveMetadata);
metadataContext.setTransitiveMetadata(Collections.unmodifiableMap(mergedTransitiveMetadata));
Map<String, String> mergedDisposableMetadata = new HashMap<>(dynamicDisposableMetadata);
metadataContext.setUpstreamDisposableMetadata(Collections.unmodifiableMap(mergedDisposableMetadata));
Map<String, String> staticDisposableMetadata = metadataContext.getFragmentContext(FRAGMENT_DISPOSABLE);
metadataContext.setDisposableMetadata(Collections.unmodifiableMap(staticDisposableMetadata));
}
if (!CollectionUtils.isEmpty(dynamicDisposableMetadata)) {
Map<String, String> mergedUpstreamDisposableMetadata = new HashMap<>(dynamicDisposableMetadata);
metadataContext.setUpstreamDisposableMetadata(Collections.unmodifiableMap(mergedUpstreamDisposableMetadata));
}
Map<String, String> staticDisposableMetadata = metadataContext.getFragmentContext(FRAGMENT_DISPOSABLE);
metadataContext.setDisposableMetadata(Collections.unmodifiableMap(staticDisposableMetadata));
MetadataContextHolder.set(metadataContext);
}

@ -22,6 +22,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import com.tencent.cloud.common.metadata.config.MetadataLocalProperties;
import com.tencent.cloud.common.spi.InstanceMetadataProvider;
@ -81,14 +82,14 @@ public class StaticMetadataManager {
private String campus;
public StaticMetadataManager(MetadataLocalProperties metadataLocalProperties,
InstanceMetadataProvider instanceMetadataProvider) {
List<InstanceMetadataProvider> instanceMetadataProviders) {
parseConfigMetadata(metadataLocalProperties);
parseEnvMetadata();
parseCustomMetadata(instanceMetadataProvider);
parseCustomMetadata(instanceMetadataProviders);
parseLocationMetadata(metadataLocalProperties, instanceMetadataProvider);
parseLocationMetadata(metadataLocalProperties, instanceMetadataProviders);
merge();
@ -182,14 +183,20 @@ public class StaticMetadataManager {
}
@SuppressWarnings("DuplicatedCode")
private void parseCustomMetadata(InstanceMetadataProvider instanceMetadataProvider) {
if (instanceMetadataProvider == null) {
private void parseCustomMetadata(List<InstanceMetadataProvider> instanceMetadataProviders) {
if (CollectionUtils.isEmpty(instanceMetadataProviders)) {
customSPIMetadata = Collections.emptyMap();
customSPITransitiveMetadata = Collections.emptyMap();
customSPIDisposableMetadata = Collections.emptyMap();
return;
}
instanceMetadataProviders.forEach(this::parseCustomMetadata);
}
@SuppressWarnings("DuplicatedCode")
private void parseCustomMetadata(InstanceMetadataProvider instanceMetadataProvider) {
// resolve all metadata
Map<String, String> allMetadata = instanceMetadataProvider.getMetadata();
if (allMetadata == null) {
@ -247,10 +254,16 @@ public class StaticMetadataManager {
}
private void parseLocationMetadata(MetadataLocalProperties metadataLocalProperties,
InstanceMetadataProvider instanceMetadataProvider) {
List<InstanceMetadataProvider> instanceMetadataProviders) {
// resolve region info
if (instanceMetadataProvider != null) {
region = instanceMetadataProvider.getRegion();
if (!CollectionUtils.isEmpty(instanceMetadataProviders)) {
Set<String> providerRegions = instanceMetadataProviders.stream().map(InstanceMetadataProvider::getRegion).filter(region -> !StringUtils.isBlank(region)).collect(Collectors.toSet());
if (!CollectionUtils.isEmpty(providerRegions)) {
if (providerRegions.size() > 1) {
throw new IllegalArgumentException("Multiple Regions Provided in InstanceMetadataProviders");
}
region = providerRegions.iterator().next();
}
}
if (StringUtils.isBlank(region)) {
region = System.getenv(ENV_METADATA_REGION);
@ -260,8 +273,14 @@ public class StaticMetadataManager {
}
// resolve zone info
if (instanceMetadataProvider != null) {
zone = instanceMetadataProvider.getZone();
if (!CollectionUtils.isEmpty(instanceMetadataProviders)) {
Set<String> providerZones = instanceMetadataProviders.stream().map(InstanceMetadataProvider::getZone).filter(zone -> !StringUtils.isBlank(zone)).collect(Collectors.toSet());
if (!CollectionUtils.isEmpty(providerZones)) {
if (providerZones.size() > 1) {
throw new IllegalArgumentException("Multiple Zones Provided in InstanceMetadataProviders");
}
zone = providerZones.iterator().next();
}
}
if (StringUtils.isBlank(zone)) {
zone = System.getenv(ENV_METADATA_ZONE);
@ -271,8 +290,14 @@ public class StaticMetadataManager {
}
// resolve campus info
if (instanceMetadataProvider != null) {
campus = instanceMetadataProvider.getCampus();
if (!CollectionUtils.isEmpty(instanceMetadataProviders)) {
Set<String> providerCampus = instanceMetadataProviders.stream().map(InstanceMetadataProvider::getCampus).filter(campus -> !StringUtils.isBlank(campus)).collect(Collectors.toSet());
if (!CollectionUtils.isEmpty(providerCampus)) {
if (providerCampus.size() > 1) {
throw new IllegalArgumentException("Multiple Campus Provided in InstanceMetadataProviders");
}
campus = providerCampus.iterator().next();
}
}
if (StringUtils.isBlank(campus)) {
campus = System.getenv(ENV_METADATA_CAMPUS);

@ -18,8 +18,12 @@
package com.tencent.cloud.common.metadata.config;
import java.util.List;
import com.tencent.cloud.common.metadata.StaticMetadataManager;
import com.tencent.cloud.common.spi.InstanceMetadataProvider;
import com.tencent.cloud.common.spi.impl.DefaultInstanceMetadataProvider;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
@ -42,10 +46,15 @@ public class MetadataAutoConfiguration {
return new MetadataLocalProperties();
}
@Bean
public InstanceMetadataProvider defaultInstanceMetadataProvider(ApplicationContextAwareUtils applicationContextAwareUtils) {
return new DefaultInstanceMetadataProvider(applicationContextAwareUtils);
}
@Bean
public StaticMetadataManager metadataManager(MetadataLocalProperties metadataLocalProperties,
@Nullable InstanceMetadataProvider instanceMetadataProvider) {
return new StaticMetadataManager(metadataLocalProperties, instanceMetadataProvider);
@Nullable List<InstanceMetadataProvider> instanceMetadataProviders) {
return new StaticMetadataManager(metadataLocalProperties, instanceMetadataProviders);
}
}

@ -0,0 +1,63 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*
*/
package com.tencent.cloud.common.spi.impl;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import com.tencent.cloud.common.spi.InstanceMetadataProvider;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
import static com.tencent.cloud.common.metadata.MetadataContext.LOCAL_NAMESPACE;
import static com.tencent.cloud.common.metadata.MetadataContext.LOCAL_SERVICE;
/**
* DefaultInstanceMetadataProvider.
* provide DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, DEFAULT_METADATA_SOURCE_SERVICE_NAME
*
* @author sean yu
*/
public class DefaultInstanceMetadataProvider implements InstanceMetadataProvider {
private final ApplicationContextAwareUtils applicationContextAwareUtils;
// ensure ApplicationContextAwareUtils init before
public DefaultInstanceMetadataProvider(ApplicationContextAwareUtils applicationContextAwareUtils) {
this.applicationContextAwareUtils = applicationContextAwareUtils;
}
@Override
public Map<String, String> getMetadata() {
return new HashMap<String, String>() {{
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, LOCAL_NAMESPACE);
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, LOCAL_SERVICE);
}};
}
@Override
public Set<String> getDisposableMetadataKeys() {
return new HashSet<>(Arrays.asList(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, DEFAULT_METADATA_SOURCE_SERVICE_NAME));
}
}

@ -93,7 +93,7 @@ public class StaticMetadataManagerTest {
when(metadataLocalProperties.getTransitive()).thenReturn(Collections.singletonList("k1"));
StaticMetadataManager metadataManager = new StaticMetadataManager(metadataLocalProperties,
new MockedMetadataProvider());
Collections.singletonList(new MockedMetadataProvider()));
Map<String, String> metadata = metadataManager.getAllCustomMetadata();
assertThat(metadata.size()).isEqualTo(3);
@ -126,7 +126,7 @@ public class StaticMetadataManagerTest {
when(metadataLocalProperties.getTransitive()).thenReturn(Collections.singletonList("k1"));
StaticMetadataManager metadataManager = new StaticMetadataManager(metadataLocalProperties,
new MockedMetadataProvider());
Collections.singletonList(new MockedMetadataProvider()));
Map<String, String> metadata = metadataManager.getMergedStaticMetadata();
assertThat(metadata.size()).isEqualTo(6);

@ -18,6 +18,8 @@
package com.tencent.cloud.common.metadata.config;
import com.tencent.cloud.common.util.ApplicationContextAwareUtils;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;
@ -33,11 +35,17 @@ import org.springframework.boot.test.context.runner.WebApplicationContextRunner;
*/
public class MetadataAutoConfigurationTest {
private final ApplicationContextRunner applicationContextRunner = new ApplicationContextRunner();
private final ApplicationContextRunner applicationContextRunner = new ApplicationContextRunner().withPropertyValues(
"spring.application.name=test"
);
private final WebApplicationContextRunner webApplicationContextRunner = new WebApplicationContextRunner();
private final WebApplicationContextRunner webApplicationContextRunner = new WebApplicationContextRunner().withPropertyValues(
"spring.application.name=test"
);
private final ReactiveWebApplicationContextRunner reactiveWebApplicationContextRunner = new ReactiveWebApplicationContextRunner();
private final ReactiveWebApplicationContextRunner reactiveWebApplicationContextRunner = new ReactiveWebApplicationContextRunner().withPropertyValues(
"spring.application.name=test"
);
/**
* No any web application.
@ -45,7 +53,7 @@ public class MetadataAutoConfigurationTest {
@Test
public void test1() {
this.applicationContextRunner
.withConfiguration(AutoConfigurations.of(MetadataAutoConfiguration.class))
.withConfiguration(AutoConfigurations.of(MetadataAutoConfiguration.class, ApplicationContextAwareUtils.class))
.run(context -> {
Assertions.assertThat(context).hasSingleBean(MetadataLocalProperties.class);
});
@ -57,7 +65,7 @@ public class MetadataAutoConfigurationTest {
@Test
public void test2() {
this.webApplicationContextRunner
.withConfiguration(AutoConfigurations.of(MetadataAutoConfiguration.class))
.withConfiguration(AutoConfigurations.of(MetadataAutoConfiguration.class, ApplicationContextAwareUtils.class))
.run(context -> {
Assertions.assertThat(context).hasSingleBean(MetadataLocalProperties.class);
});
@ -69,7 +77,7 @@ public class MetadataAutoConfigurationTest {
@Test
public void test3() {
this.reactiveWebApplicationContextRunner
.withConfiguration(AutoConfigurations.of(MetadataAutoConfiguration.class))
.withConfiguration(AutoConfigurations.of(MetadataAutoConfiguration.class, ApplicationContextAwareUtils.class))
.run(context -> {
Assertions.assertThat(context).hasSingleBean(MetadataLocalProperties.class);
});

@ -18,6 +18,11 @@
<artifactId>spring-boot-starter-web</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-webflux</artifactId>
</dependency>
<dependency>
<groupId>com.tencent.cloud</groupId>
<artifactId>spring-cloud-starter-tencent-polaris-discovery</artifactId>

@ -17,15 +17,22 @@
package com.tencent.cloud.ratelimit.example.service.callee;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
@ -33,6 +40,8 @@ import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.HttpClientErrorException.TooManyRequests;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
/**
* Rate limit controller.
@ -49,6 +58,8 @@ public class BusinessController {
private final AtomicLong lastTimestamp = new AtomicLong(0);
@Autowired
private RestTemplate restTemplate;
@Autowired
private WebClient.Builder webClientBuilder;
@Value("${spring.application.name}")
private String appName;
@ -62,6 +73,44 @@ public class BusinessController {
return "hello world for ratelimit service " + index.incrementAndGet();
}
@GetMapping("/info/webclient")
public Mono<String> infoWebClient() {
return Mono.just("hello world for ratelimit service " + index.incrementAndGet());
}
@GetMapping("/invoke/webclient")
public String invokeInfoWebClient() throws InterruptedException, ExecutionException {
StringBuffer builder = new StringBuffer();
WebClient webClient = webClientBuilder.baseUrl("http://" + appName).build();
List<Mono<String>> monoList = new ArrayList<>();
for (int i = 0; i < 30; i++) {
Mono<String> response = webClient.get()
.uri(uriBuilder -> uriBuilder
.path("/business/info/webclient")
.queryParam("yyy", "yyy")
.build()
)
.header("xxx", "xxx")
.retrieve()
.bodyToMono(String.class)
.doOnSuccess(s -> builder.append(s + "\n"))
.doOnError(e -> {
if (e instanceof WebClientResponseException) {
if (((WebClientResponseException) e).getRawStatusCode() == 429) {
builder.append("TooManyRequests ").append(index.incrementAndGet() + "\n");
}
}
})
.onErrorReturn("");
monoList.add(response);
}
for (Mono<String> mono : monoList) {
mono.toFuture().get();
}
index.set(0);
return builder.toString();
}
/**
* Get information 30 times per 1 second.
*
@ -75,19 +124,31 @@ public class BusinessController {
for (int i = 0; i < 30; i++) {
new Thread(() -> {
try {
ResponseEntity<String> entity = restTemplate.getForEntity("http://" + appName + "/business/info",
String.class);
HttpHeaders httpHeaders = new HttpHeaders();
httpHeaders.add("xxx", "xxx");
ResponseEntity<String> entity = restTemplate.exchange(
"http://" + appName + "/business/info?yyy={yyy}",
HttpMethod.GET,
new HttpEntity<>(httpHeaders),
String.class,
"yyy"
);
builder.append(entity.getBody() + "\n");
}
catch (RestClientException e) {
if (e instanceof TooManyRequests) {
builder.append("TooManyRequests " + index.incrementAndGet() + "\n");
builder.append("TooManyRequests ").append(index.incrementAndGet() + "\n");
}
else {
throw e;
}
}
count.countDown();
catch (Exception e) {
e.printStackTrace();
}
finally {
count.countDown();
}
}).start();
}
count.await();

@ -0,0 +1,44 @@
/*
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
*
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
*
* Licensed under the BSD 3-Clause License (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://opensource.org/licenses/BSD-3-Clause
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/
package com.tencent.cloud.ratelimit.example.service.callee;
import java.util.HashMap;
import java.util.Map;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
/**
* resolver custom label from request.
*
* @author sean yu
*/
@Component
public class CustomLabelResolverReactive implements PolarisRateLimiterLabelReactiveResolver {
@Override
public Map<String, String> resolve(ServerWebExchange exchange) {
// rate limit by some request params. such as query params, headers ..
Map<String, String> labels = new HashMap<>();
labels.put("user", "zhangsan");
return labels;
}
}

@ -22,6 +22,7 @@ import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.cloud.client.loadbalancer.LoadBalanced;
import org.springframework.context.annotation.Bean;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
/**
* Rate limit application.
@ -40,4 +41,11 @@ public class RateLimitCalleeService {
public RestTemplate restTemplate() {
return new RestTemplate();
}
@LoadBalanced
@Bean
WebClient.Builder webClientBuilder() {
return WebClient.builder();
}
}

@ -50,7 +50,7 @@ public class RpcEnhancementAutoConfigurationTest {
RpcEnhancementAutoConfiguration.class,
PolarisRestTemplateAutoConfigurationTester.class,
FeignLoadBalancerAutoConfiguration.class))
.withPropertyValues("spring.cloud.polaris.circuitbreaker.enabled=true");
.withPropertyValues("spring.cloud.polaris.circuitbreaker.enabled=true", "spring.application.name=test");
@Test
public void testDefaultInitialization() {

@ -38,7 +38,7 @@ import static org.springframework.http.HttpStatus.Series.SERVER_ERROR;
* @author Haotian Zhang
*/
@ExtendWith(SpringExtension.class)
@SpringBootTest(classes = RpcEnhancementReporterPropertiesTest.TestApplication.class)
@SpringBootTest(classes = RpcEnhancementReporterPropertiesTest.TestApplication.class, properties = "spring.application.name=test")
@ActiveProfiles("test")
public class RpcEnhancementReporterPropertiesTest {

@ -41,7 +41,8 @@ public class StatConfigModifierTest {
.withPropertyValues("spring.cloud.polaris.stat.enabled=true")
.withPropertyValues("spring.cloud.polaris.stat.host=127.0.0.1")
.withPropertyValues("spring.cloud.polaris.stat.port=20000")
.withPropertyValues("spring.cloud.polaris.stat.path=/xxx");
.withPropertyValues("spring.cloud.polaris.stat.path=/xxx")
.withPropertyValues("spring.application.name=test");
private final ApplicationContextRunner pushContextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(TestApplication.class))
@ -49,12 +50,14 @@ public class StatConfigModifierTest {
.withPropertyValues("spring.cloud.polaris.stat.enabled=true")
.withPropertyValues("spring.cloud.polaris.stat.pushgateway.enabled=true")
.withPropertyValues("spring.cloud.polaris.stat.pushgateway.address=127.0.0.1:9091")
.withPropertyValues("spring.cloud.polaris.stat.pushgateway.push-interval=1000");
.withPropertyValues("spring.cloud.polaris.stat.pushgateway.push-interval=1000")
.withPropertyValues("spring.application.name=test");
private final ApplicationContextRunner disabledContextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(TestApplication.class))
.withPropertyValues("spring.cloud.polaris.enabled=true")
.withPropertyValues("spring.cloud.polaris.stat.enabled=false");
.withPropertyValues("spring.cloud.polaris.stat.enabled=false")
.withPropertyValues("spring.application.name=test");
@Test
void testPull() {

Loading…
Cancel
Save