diff --git a/CHANGELOG.md b/CHANGELOG.md index 35f38c02..52e75767 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,7 @@ # Change Log --- +- [feat:add custom label resolver spi for rate limit](https://github.com/Tencent/spring-cloud-tencent/pull/107) - [feat:fix discovery weight param not set to register request bug](https://github.com/Tencent/spring-cloud-tencent/pull/104) - [Bugfix: fix causing cpu 100% when set ScheduledThreadPoolExecutor corePoolSize=0](https://github.com/Tencent/spring-cloud-tencent/pull/101) + diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/RateLimitConfiguration.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/RateLimitConfiguration.java index 7c87a300..39499e53 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/RateLimitConfiguration.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/RateLimitConfiguration.java @@ -13,6 +13,7 @@ * 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.config; @@ -20,6 +21,8 @@ package com.tencent.cloud.polaris.ratelimit.config; 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.spi.PolarisRateLimiterLabelReactiveResolver; +import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver; import com.tencent.polaris.client.api.SDKContext; import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.factory.LimitAPIFactory; @@ -30,6 +33,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplicat import org.springframework.boot.web.servlet.FilterRegistrationBean; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.lang.Nullable; import static javax.servlet.DispatcherType.ASYNC; import static javax.servlet.DispatcherType.ERROR; @@ -59,8 +63,9 @@ public class RateLimitConfiguration { @Bean @ConditionalOnMissingBean - public QuotaCheckServletFilter quotaCheckFilter(LimitAPI limitAPI) { - return new QuotaCheckServletFilter(limitAPI); + public QuotaCheckServletFilter quotaCheckFilter(LimitAPI limitAPI, + @Nullable PolarisRateLimiterLabelServletResolver labelResolver) { + return new QuotaCheckServletFilter(limitAPI, labelResolver); } @Bean @@ -84,8 +89,9 @@ public class RateLimitConfiguration { static class MetadataReactiveFilterConfig { @Bean - public QuotaCheckReactiveFilter quotaCheckReactiveFilter(LimitAPI limitAPI) { - return new QuotaCheckReactiveFilter(limitAPI); + public QuotaCheckReactiveFilter quotaCheckReactiveFilter(LimitAPI limitAPI, + @Nullable PolarisRateLimiterLabelReactiveResolver labelResolver) { + return new QuotaCheckReactiveFilter(limitAPI, labelResolver); } } diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/constant/RateLimitConstant.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/constant/RateLimitConstant.java index 5c96b864..81413452 100644 --- a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/constant/RateLimitConstant.java +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/constant/RateLimitConstant.java @@ -37,6 +37,11 @@ public final class RateLimitConstant { */ public static String QUOTA_LIMITED_INFO = "The request is deny by rate limit because the throttling threshold is reached"; + /** + * The build in label method. + */ + public static String LABEL_METHOD = "method"; + private RateLimitConstant() { } 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 ad6a0049..f838e1a7 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 @@ -13,6 +13,7 @@ * 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.filter; @@ -23,6 +24,7 @@ import java.util.Map; import com.tencent.cloud.common.metadata.MetadataContext; import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant; +import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver; import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils; import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse; @@ -37,10 +39,13 @@ import org.springframework.core.io.buffer.DataBuffer; import org.springframework.http.HttpStatus; import org.springframework.http.MediaType; import org.springframework.http.server.reactive.ServerHttpResponse; +import org.springframework.util.CollectionUtils; 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. * @@ -52,8 +57,12 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered { private final LimitAPI limitAPI; - public QuotaCheckReactiveFilter(LimitAPI limitAPI) { + private final PolarisRateLimiterLabelReactiveResolver labelResolver; + + public QuotaCheckReactiveFilter(LimitAPI limitAPI, + PolarisRateLimiterLabelReactiveResolver labelResolver) { this.limitAPI = limitAPI; + this.labelResolver = labelResolver; } @Override @@ -65,23 +74,37 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered { public Mono filter(ServerWebExchange exchange, WebFilterChain chain) { String localNamespace = MetadataContext.LOCAL_NAMESPACE; String localService = MetadataContext.LOCAL_SERVICE; - String method = exchange.getRequest().getURI().getPath(); - Map labels = null; - if (StringUtils.isNotBlank(method)) { - labels = new HashMap<>(); - labels.put("method", method); + + Map labels = new HashMap<>(); + + // add build in labels + String path = exchange.getRequest().getURI().getPath(); + if (StringUtils.isNotBlank(path)) { + labels.put(LABEL_METHOD, path); + } + + // add custom labels + if (labelResolver != null) { + try { + Map customLabels = labelResolver.resolve(exchange); + if (!CollectionUtils.isEmpty(customLabels)) { + labels.putAll(customLabels); + } + } catch (Throwable e) { + LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e); + } } try { - QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, localNamespace, localService, 1, labels, - null); + QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, + localNamespace, localService, 1, labels, null); + if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) { ServerHttpResponse response = exchange.getResponse(); response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS); response.getHeaders().setContentType(MediaType.APPLICATION_JSON); - DataBuffer dataBuffer = response.bufferFactory().allocateBuffer() - .write((RateLimitConstant.QUOTA_LIMITED_INFO + quotaResponse.getInfo()) - .getBytes(StandardCharsets.UTF_8)); + DataBuffer dataBuffer = response.bufferFactory().allocateBuffer().write( + RateLimitConstant.QUOTA_LIMITED_INFO.getBytes(StandardCharsets.UTF_8)); return response.writeWith(Mono.just(dataBuffer)); } } 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 fb386918..43a98faf 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 @@ -29,6 +29,7 @@ import javax.servlet.http.HttpServletResponse; import com.tencent.cloud.common.metadata.MetadataContext; import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant; +import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver; import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils; import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse; @@ -38,8 +39,10 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.core.annotation.Order; +import org.springframework.util.CollectionUtils; import org.springframework.web.filter.OncePerRequestFilter; +import static com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant.LABEL_METHOD; import static org.springframework.http.HttpStatus.TOO_MANY_REQUESTS; /** @@ -53,9 +56,11 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter { private static final Logger LOG = LoggerFactory.getLogger(QuotaCheckServletFilter.class); private final LimitAPI limitAPI; + private final PolarisRateLimiterLabelServletResolver labelResolver; - public QuotaCheckServletFilter(LimitAPI limitAPI) { + public QuotaCheckServletFilter(LimitAPI limitAPI, PolarisRateLimiterLabelServletResolver labelResolver) { this.limitAPI = limitAPI; + this.labelResolver = labelResolver; } @Override @@ -63,11 +68,26 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter { throws ServletException, IOException { String localNamespace = MetadataContext.LOCAL_NAMESPACE; String localService = MetadataContext.LOCAL_SERVICE; - String method = request.getRequestURI(); - Map labels = null; - if (StringUtils.isNotBlank(method)) { - labels = new HashMap<>(); - labels.put("method", method); + + Map labels = new HashMap<>(); + + // add build in labels + String path = request.getRequestURI(); + + if (StringUtils.isNotBlank(path)) { + labels.put(LABEL_METHOD, path); + } + + // add custom labels + if (labelResolver != null) { + try { + Map customLabels = labelResolver.resolve(request); + if (!CollectionUtils.isEmpty(customLabels)) { + labels.putAll(customLabels); + } + } catch (Throwable e) { + LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e); + } } try { diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/spi/PolarisRateLimiterLabelReactiveResolver.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/spi/PolarisRateLimiterLabelReactiveResolver.java new file mode 100644 index 00000000..790c7c0d --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/spi/PolarisRateLimiterLabelReactiveResolver.java @@ -0,0 +1,40 @@ +/* + * 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.spi; + +import java.util.Map; + +import org.springframework.web.server.ServerWebExchange; + +/** + * Resolve custom label from request. The label used for rate limit params. + * + * @author lepdou 2022-03-31 + */ +public interface PolarisRateLimiterLabelReactiveResolver { + + /** + * Resolve custom label from request. + * + * @param exchange the http request + * @return resolved labels + */ + Map resolve(ServerWebExchange exchange); + +} diff --git a/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/spi/PolarisRateLimiterLabelServletResolver.java b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/spi/PolarisRateLimiterLabelServletResolver.java new file mode 100644 index 00000000..8fe85909 --- /dev/null +++ b/spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/spi/PolarisRateLimiterLabelServletResolver.java @@ -0,0 +1,40 @@ +/* + * 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.spi; + +import java.util.Map; + +import javax.servlet.http.HttpServletRequest; + +/** + * Resolve custom label from request. The label used for rate limit params. + * + * @author lepdou 2022-03-31 + */ +public interface PolarisRateLimiterLabelServletResolver { + + /** + * Resolve custom label from request. + * + * @param request the http request + * @return resolved labels + */ + Map resolve(HttpServletRequest request); + +} diff --git a/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolver.java b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolver.java new file mode 100644 index 00000000..837de09c --- /dev/null +++ b/spring-cloud-tencent-examples/polaris-ratelimit-example/ratelimit-callee-service/src/main/java/com/tencent/cloud/ratelimit/example/service/callee/CustomLabelResolver.java @@ -0,0 +1,47 @@ +/* + * 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 javax.servlet.http.HttpServletRequest; + +import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver; + +import org.springframework.stereotype.Component; + +/** + * resolver custom label from request. + * + *@author lepdou 2022-03-31 + */ +@Component +public class CustomLabelResolver implements PolarisRateLimiterLabelServletResolver { + + @Override + public Map resolve(HttpServletRequest request) { + //rate limit by some request params. such as query params, headers .. + + Map labels = new HashMap<>(); + labels.put("user", "zhangsan"); + + return labels; + } +}