From ad9b485bbe0d63482f77b75dd297011d6abfa534 Mon Sep 17 00:00:00 2001 From: "Shanyou Yu (Sean Yu)" Date: Mon, 27 Mar 2023 11:49:17 +0800 Subject: [PATCH] feature: add ratelimit provider info and refactor ratelimit use arguments (#903) --- CHANGELOG.md | 3 +- .../MetadataTransferAutoConfiguration.java | 24 +++ .../EncodeTransferMedataWebClientFilter.java | 90 ++++++++++ ...MetadataTransferAutoConfigurationTest.java | 2 + ...codeTransferMedataWebClientFilterTest.java | 84 +++++++++ .../PolarisRouterAutoConfigurationTest.java | 2 +- .../cloud/polaris/util/OkHttpUtilTest.java | 2 +- .../pom.xml | 5 + .../ratelimit/RateLimitRuleLabelResolver.java | 1 + .../PolarisRateLimitAutoConfiguration.java | 41 +++-- .../filter/QuotaCheckReactiveFilter.java | 61 +------ .../filter/QuotaCheckServletFilter.java | 61 +------ ...RateLimitRuleArgumentReactiveResolver.java | 123 ++++++++++++++ .../RateLimitRuleArgumentServletResolver.java | 123 ++++++++++++++ .../ratelimit/utils/QuotaCheckUtils.java | 22 +++ ...PolarisRateLimitAutoConfigurationTest.java | 12 +- .../PolarisRateLimitPropertiesTest.java | 2 +- .../filter/QuotaCheckReactiveFilterTest.java | 122 +++++-------- .../filter/QuotaCheckServletFilterTest.java | 160 +++++++----------- ...LimitRuleArgumentReactiveResolverTest.java | 144 ++++++++++++++++ ...eLimitRuleArgumentServletResolverTest.java | 135 +++++++++++++++ .../ratelimit/utils/QuotaCheckUtilsTest.java | 42 ++++- .../src/test/resources/ratelimit.json | 95 +++++++++++ .../config/FeignAutoConfigurationTest.java | 5 +- .../config/RouterAutoConfigurationTests.java | 5 +- .../common/constant/MetadataConstant.java | 17 +- .../metadata/MetadataContextHolder.java | 12 +- .../metadata/StaticMetadataManager.java | 49 ++++-- .../config/MetadataAutoConfiguration.java | 13 +- .../impl/DefaultInstanceMetadataProvider.java | 63 +++++++ .../metadata/StaticMetadataManagerTest.java | 4 +- .../config/MetadataAutoConfigurationTest.java | 20 ++- .../ratelimit-callee-service/pom.xml | 5 + .../service/callee/BusinessController.java | 78 ++++++++- .../callee/CustomLabelResolverReactive.java | 44 +++++ .../callee/RateLimitCalleeService.java | 8 + .../RpcEnhancementAutoConfigurationTest.java | 2 +- .../RpcEnhancementReporterPropertiesTest.java | 2 +- .../stat/config/StatConfigModifierTest.java | 9 +- 39 files changed, 1328 insertions(+), 364 deletions(-) create mode 100644 spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilter.java create mode 100644 spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilterTest.java create mode 100644 spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolver.java create mode 100644 spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolver.java create mode 100644 spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolverTest.java create mode 100644 spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolverTest.java create mode 100644 spring-cloud-starter-tencent-polaris-ratelimit/src/test/resources/ratelimit.json create mode 100644 spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/spi/impl/DefaultInstanceMetadataProvider.java create mode 100644 spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolverReactive.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 3cae1b0e6..002c72b14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,9 @@ # Change Log --- -- [fix:fix log feign response stream close bug.](https://github.com/Tencent/spring-cloud-tencent/pull/897) +- [fix: fix log feign response stream close bug.](https://github.com/Tencent/spring-cloud-tencent/pull/897) - [fix:remove the secondary report.](https://github.com/Tencent/spring-cloud-tencent/pull/899) +- [feature: add ratelimit provider info and refactor ratelimit use arguments.](https://github.com/Tencent/spring-cloud-tencent/pull/903) - [fix:optimize instance circuit beaker.](https://github.com/Tencent/spring-cloud-tencent/pull/909) - [fix:optimize multi service registration and discovery.](https://github.com/Tencent/spring-cloud-tencent/pull/912) - [feature: improve circuit breaker usage.](https://github.com/Tencent/spring-cloud-tencent/pull/913) diff --git a/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfiguration.java b/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfiguration.java index 7b2aa8f9b..f922815f3 100644 --- a/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfiguration.java +++ b/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfiguration.java @@ -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; @@ -140,4 +142,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 webClientBuilder = Collections.emptyList(); + + @Bean + public EncodeTransferMedataWebClientFilter encodeTransferMedataWebClientFilter() { + return new EncodeTransferMedataWebClientFilter(); + } + + @Bean + public SmartInitializingSingleton addEncodeTransferMetadataFilterForWebClient(EncodeTransferMedataWebClientFilter filter) { + return () -> webClientBuilder.forEach(webClient -> { + webClient.filter(filter); + }); + } + } } diff --git a/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilter.java b/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilter.java new file mode 100644 index 000000000..3547add0e --- /dev/null +++ b/spring-cloud-starter-tencent-metadata-transfer/src/main/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilter.java @@ -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 filter(ClientRequest clientRequest, ExchangeFunction next) { + MetadataContext metadataContext = MetadataContextHolder.get(); + Map customMetadata = metadataContext.getCustomMetadata(); + Map disposableMetadata = metadataContext.getDisposableMetadata(); + Map 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 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 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); + } + } + } + +} diff --git a/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfigurationTest.java b/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfigurationTest.java index 7f3164e3e..12073bc9e 100644 --- a/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfigurationTest.java +++ b/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/config/MetadataTransferAutoConfigurationTest.java @@ -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); }); } diff --git a/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilterTest.java b/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilterTest.java new file mode 100644 index 000000000..ec44eec32 --- /dev/null +++ b/spring-cloud-starter-tencent-metadata-transfer/src/test/java/com/tencent/cloud/metadata/core/EncodeTransferMedataWebClientFilterTest.java @@ -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"); + } + } +} diff --git a/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/loadbalancer/PolarisRouterAutoConfigurationTest.java b/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/loadbalancer/PolarisRouterAutoConfigurationTest.java index 123343e72..9e9fc863d 100644 --- a/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/loadbalancer/PolarisRouterAutoConfigurationTest.java +++ b/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/loadbalancer/PolarisRouterAutoConfigurationTest.java @@ -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( diff --git a/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/util/OkHttpUtilTest.java b/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/util/OkHttpUtilTest.java index 7590b8b1e..b718e88e9 100644 --- a/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/util/OkHttpUtilTest.java +++ b/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/util/OkHttpUtilTest.java @@ -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 diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/pom.xml b/spring-cloud-starter-tencent-polaris-ratelimit/pom.xml index 5df5c0183..1c866e10a 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/pom.xml +++ b/spring-cloud-starter-tencent-polaris-ratelimit/pom.xml @@ -19,6 +19,11 @@ com.tencent.cloud spring-cloud-tencent-rpc-enhancement + + + com.tencent.cloud + spring-cloud-starter-tencent-metadata-transfer + diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/RateLimitRuleLabelResolver.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/RateLimitRuleLabelResolver.java index fed34dcc8..f51b6f0b8 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/RateLimitRuleLabelResolver.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/RateLimitRuleLabelResolver.java @@ -36,6 +36,7 @@ import org.springframework.util.CollectionUtils; * *@author lepdou 2022-05-13 */ +@Deprecated public class RateLimitRuleLabelResolver { private final ServiceRuleManager serviceRuleManager; diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfiguration.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfiguration.java index b8b79b7a9..d33c0e604 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfiguration.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfiguration.java @@ -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); } } } diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilter.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilter.java index 7fff409d7..db0075a90 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilter.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilter.java @@ -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 labels = getRequestLabels(exchange, localNamespace, localService); + Set 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); + 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 getRequestLabels(ServerWebExchange exchange, String localNamespace, String localService) { - Map 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 expressionLabels = getRuleExpressionLabels(exchange, localNamespace, localService); - labels.putAll(expressionLabels); - - // add custom labels - Map customResolvedLabels = getCustomResolvedLabels(exchange); - labels.putAll(customResolvedLabels); - - return labels; - } - - private Map 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 getRuleExpressionLabels(ServerWebExchange exchange, String namespace, String service) { - Set expressionLabels = rateLimitRuleLabelResolver.getExpressionLabelKeys(namespace, service); - return SpringWebExpressionLabelUtils.resolve(exchange, expressionLabels); - } } diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilter.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilter.java index 434046434..1bc3a5203 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilter.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilter.java @@ -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 labels = getRequestLabels(request, localNamespace, localService); + Set 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 getRequestLabels(HttpServletRequest request, String localNamespace, String localService) { - Map 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 expressionLabels = getRuleExpressionLabels(request, localNamespace, localService); - labels.putAll(expressionLabels); - - // add custom resolved labels - Map customLabels = getCustomResolvedLabels(request); - labels.putAll(customLabels); - - return labels; - } - - private Map 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 getRuleExpressionLabels(HttpServletRequest request, String namespace, String service) { - Set expressionLabels = rateLimitRuleLabelResolver.getExpressionLabelKeys(namespace, service); - return ServletExpressionLabelUtils.resolve(request, expressionLabels); - } } diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolver.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolver.java new file mode 100644 index 000000000..724a5eab6 --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolver.java @@ -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 getArguments(ServerWebExchange request, String namespace, String service) { + RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service); + if (rateLimitRule == null) { + return Collections.emptySet(); + } + List 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 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(); + } +} diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolver.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolver.java new file mode 100644 index 000000000..e384d0333 --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolver.java @@ -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 getArguments(HttpServletRequest request, String namespace, String service) { + RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service); + if (rateLimitRule == null) { + return Collections.emptySet(); + } + List 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 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(); + } + +} diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtils.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtils.java index faac3490d..3adcee5e4 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtils.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtils.java @@ -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 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 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")); + } + } } diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfigurationTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfigurationTest.java index ac532a8bc..372d2c8c2 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfigurationTest.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitAutoConfigurationTest.java @@ -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; @@ -54,7 +55,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); }); } @@ -90,7 +93,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); diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitPropertiesTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitPropertiesTest.java index e36b9d530..3b27fe84a 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitPropertiesTest.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/config/PolarisRateLimitPropertiesTest.java @@ -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 { diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilterTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilterTest.java index e707cb475..5b10a90b4 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilterTest.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilterTest.java @@ -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 mockedApplicationContextAwareUtils; - private static MockedStatic 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("

RejectRequestTips提示消息

"); + 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 result = (Map) 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) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange); - assertThat(result.size()).isEqualTo(0); - - // labelResolver == null - quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, null, null, null, null); - result = (Map) 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 diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilterTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilterTest.java index 411d51bf7..b5f64f882 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilterTest.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckServletFilterTest.java @@ -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 mockedApplicationContextAwareUtils; - private static MockedStatic 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("

RejectRequestTips提示消息

"); 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 result = (Map) 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) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request); - assertThat(result.size()).isEqualTo(0); - - // labelResolver == null - quotaCheckServletFilter = new QuotaCheckServletFilter(null, null, null, null, null); - result = (Map) 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,33 +158,38 @@ 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("

RejectRequestTips提示消息

"); // 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); @@ -251,31 +204,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); diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolverTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolverTest.java new file mode 100644 index 000000000..1648a9618 --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentReactiveResolverTest.java @@ -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() {{ + put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE); + put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE); + }}); + MetadataContextHolder.set(metadataContext); + Set arguments = rateLimitRuleArgumentReactiveResolver1.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE); + Set 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 { + } + +} diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolverTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolverTest.java new file mode 100644 index 000000000..1a8b49ed7 --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/resolver/RateLimitRuleArgumentServletResolverTest.java @@ -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() {{ + put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE); + put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE); + }}); + MetadataContextHolder.set(metadataContext); + Set arguments = rateLimitRuleArgumentServletResolver1.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE); + Set 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 { + } +} diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtilsTest.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtilsTest.java index 91e361a8f..ad2691f4d 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtilsTest.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/java/com/tencent/cloud/polaris/ratelimit/utils/QuotaCheckUtilsTest.java @@ -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"); diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/test/resources/ratelimit.json b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/resources/ratelimit.json new file mode 100644 index 000000000..92908114d --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/test/resources/ratelimit.json @@ -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 +} \ No newline at end of file diff --git a/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/FeignAutoConfigurationTest.java b/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/FeignAutoConfigurationTest.java index 30f37a15f..774ca3acc 100644 --- a/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/FeignAutoConfigurationTest.java +++ b/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/FeignAutoConfigurationTest.java @@ -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() { diff --git a/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/RouterAutoConfigurationTests.java b/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/RouterAutoConfigurationTests.java index 5416a5644..3bf557430 100644 --- a/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/RouterAutoConfigurationTests.java +++ b/spring-cloud-starter-tencent-polaris-router/src/test/java/com/tencent/cloud/polaris/router/config/RouterAutoConfigurationTests.java @@ -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() { diff --git a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/constant/MetadataConstant.java b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/constant/MetadataConstant.java index 1a6cb57d6..c3774baf1 100644 --- a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/constant/MetadataConstant.java +++ b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/constant/MetadataConstant.java @@ -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"; + + } + } diff --git a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/MetadataContextHolder.java b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/MetadataContextHolder.java index fa325879a..f6307ae32 100644 --- a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/MetadataContextHolder.java +++ b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/MetadataContextHolder.java @@ -135,13 +135,13 @@ public final class MetadataContextHolder { mergedTransitiveMetadata.putAll(staticTransitiveMetadata); mergedTransitiveMetadata.putAll(dynamicTransitiveMetadata); metadataContext.setTransitiveMetadata(Collections.unmodifiableMap(mergedTransitiveMetadata)); - - Map mergedDisposableMetadata = new HashMap<>(dynamicDisposableMetadata); - metadataContext.setUpstreamDisposableMetadata(Collections.unmodifiableMap(mergedDisposableMetadata)); - - Map staticDisposableMetadata = metadataContext.getDisposableMetadata(); - metadataContext.setDisposableMetadata(Collections.unmodifiableMap(staticDisposableMetadata)); } + if (!CollectionUtils.isEmpty(dynamicDisposableMetadata)) { + Map mergedUpstreamDisposableMetadata = new HashMap<>(dynamicDisposableMetadata); + metadataContext.setUpstreamDisposableMetadata(Collections.unmodifiableMap(mergedUpstreamDisposableMetadata)); + } + Map staticDisposableMetadata = metadataContext.getDisposableMetadata(); + metadataContext.setDisposableMetadata(Collections.unmodifiableMap(staticDisposableMetadata)); MetadataContextHolder.set(metadataContext); } diff --git a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/StaticMetadataManager.java b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/StaticMetadataManager.java index 63ea080fb..3bf69dacd 100644 --- a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/StaticMetadataManager.java +++ b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/StaticMetadataManager.java @@ -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 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 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 allMetadata = instanceMetadataProvider.getMetadata(); if (allMetadata == null) { @@ -247,10 +254,16 @@ public class StaticMetadataManager { } private void parseLocationMetadata(MetadataLocalProperties metadataLocalProperties, - InstanceMetadataProvider instanceMetadataProvider) { + List instanceMetadataProviders) { // resolve region info - if (instanceMetadataProvider != null) { - region = instanceMetadataProvider.getRegion(); + if (!CollectionUtils.isEmpty(instanceMetadataProviders)) { + Set 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 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 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); diff --git a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfiguration.java b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfiguration.java index 060bed30d..a5a02b67b 100644 --- a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfiguration.java +++ b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfiguration.java @@ -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 instanceMetadataProviders) { + return new StaticMetadataManager(metadataLocalProperties, instanceMetadataProviders); } } diff --git a/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/spi/impl/DefaultInstanceMetadataProvider.java b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/spi/impl/DefaultInstanceMetadataProvider.java new file mode 100644 index 000000000..e4a99636c --- /dev/null +++ b/spring-cloud-tencent-commons/src/main/java/com/tencent/cloud/common/spi/impl/DefaultInstanceMetadataProvider.java @@ -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 getMetadata() { + return new HashMap() {{ + put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, LOCAL_NAMESPACE); + put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, LOCAL_SERVICE); + }}; + } + + @Override + public Set getDisposableMetadataKeys() { + return new HashSet<>(Arrays.asList(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, DEFAULT_METADATA_SOURCE_SERVICE_NAME)); + } + +} diff --git a/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/StaticMetadataManagerTest.java b/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/StaticMetadataManagerTest.java index 2bd773e98..8dbda4cfe 100644 --- a/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/StaticMetadataManagerTest.java +++ b/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/StaticMetadataManagerTest.java @@ -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 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 metadata = metadataManager.getMergedStaticMetadata(); assertThat(metadata.size()).isEqualTo(6); diff --git a/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfigurationTest.java b/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfigurationTest.java index 83143d6d5..c5cda081c 100644 --- a/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfigurationTest.java +++ b/spring-cloud-tencent-commons/src/test/java/com/tencent/cloud/common/metadata/config/MetadataAutoConfigurationTest.java @@ -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); }); diff --git a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/pom.xml b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/pom.xml index 4bbda499b..613cf24f7 100644 --- a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/pom.xml +++ b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/pom.xml @@ -18,6 +18,11 @@ spring-boot-starter-web + + org.springframework.boot + spring-boot-starter-webflux + + com.tencent.cloud spring-cloud-starter-tencent-polaris-discovery diff --git a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/BusinessController.java b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/BusinessController.java index 11d31fd58..fc8a51f3f 100644 --- a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/BusinessController.java +++ b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/BusinessController.java @@ -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; @@ -56,11 +67,55 @@ public class BusinessController { * Get information. * @return information */ - @RequestMapping("/info") + @GetMapping("/info") public String info() { return "hello world for ratelimit service " + index.incrementAndGet(); } + @GetMapping("/info/webclient") + public Mono 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> monoList = new ArrayList<>(); + for (int i = 0; i < 30; i++) { + Mono 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 mono : monoList) { + mono.toFuture().get(); + } + index.set(0); + return builder.toString(); + } + + /** + * Get information 30 times per 1 second. + * + * @return result of 30 calls. + * @throws InterruptedException exception + */ @GetMapping("/invoke") public String invokeInfo() throws InterruptedException { StringBuffer builder = new StringBuffer(); @@ -68,19 +123,31 @@ public class BusinessController { for (int i = 0; i < 30; i++) { new Thread(() -> { try { - ResponseEntity entity = restTemplate.getForEntity("http://" + appName + "/business/info", - String.class); + HttpHeaders httpHeaders = new HttpHeaders(); + httpHeaders.add("xxx", "xxx"); + ResponseEntity 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(); @@ -90,6 +157,7 @@ public class BusinessController { /** * Get information with unirate. + * * @return information */ @GetMapping("/unirate") diff --git a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolverReactive.java b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolverReactive.java new file mode 100644 index 000000000..eb6e01b03 --- /dev/null +++ b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolverReactive.java @@ -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 resolve(ServerWebExchange exchange) { + // rate limit by some request params. such as query params, headers .. + + Map labels = new HashMap<>(); + labels.put("user", "zhangsan"); + + return labels; + } +} diff --git a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/RateLimitCalleeService.java b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/RateLimitCalleeService.java index 12a9270cc..c8a49d2a7 100644 --- a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/RateLimitCalleeService.java +++ b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/RateLimitCalleeService.java @@ -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(); + } } diff --git a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java index 7be8b8b60..199fd36b7 100644 --- a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java +++ b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java @@ -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() { diff --git a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementReporterPropertiesTest.java b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementReporterPropertiesTest.java index 48db48072..111c031c5 100644 --- a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementReporterPropertiesTest.java +++ b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementReporterPropertiesTest.java @@ -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 { diff --git a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/stat/config/StatConfigModifierTest.java b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/stat/config/StatConfigModifierTest.java index ba55b6f2e..4ef678597 100644 --- a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/stat/config/StatConfigModifierTest.java +++ b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/stat/config/StatConfigModifierTest.java @@ -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() {