add custom label resolver spi for rate limit

CHANGELOG.md

@@ -2,4 +2,4 @@
 ---
 
 - [Bugfix: fix router bug and add router example](https://github.com/Tencent/spring-cloud-tencent/pull/89)
-
+- [feat:add custom label resolver spi for rate limit](https://github.com/Tencent/spring-cloud-tencent/pull/105)

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;
@@ -60,8 +64,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
@@ -85,8 +90,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);
 		}
 
 	}

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() {
 	}
 

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.
  *
@@ -53,8 +58,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
@@ -66,23 +75,37 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
 	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
 		String localNamespace = MetadataContext.LOCAL_NAMESPACE;
 		String localService = MetadataContext.LOCAL_SERVICE;
-		String method = exchange.getRequest().getURI().getPath();
-		Map<String, String> labels = null;
-		if (StringUtils.isNotBlank(method)) {
-			labels = new HashMap<>();
-			labels.put("method", method);
+
+		Map<String, String> labels = new HashMap<>();
+
+		// add build in labels
+		String path = exchange.getRequest().getURI().getPath();
+		if (StringUtils.isNotBlank(path)) {
+			labels.put(LABEL_METHOD, path);
+		}
+
+		// add custom labels
+		if (labelResolver != null) {
+			try {
+				Map<String, String> 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);
+
 			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));
+						RateLimitConstant.QUOTA_LIMITED_INFO.getBytes(StandardCharsets.UTF_8));
 				return response.writeWith(Mono.just(dataBuffer));
 			}
 		}

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;
 
 /**
@@ -54,9 +57,11 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
 			.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
@@ -65,11 +70,26 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
 			throws ServletException, IOException {
 		String localNamespace = MetadataContext.LOCAL_NAMESPACE;
 		String localService = MetadataContext.LOCAL_SERVICE;
-		String method = request.getRequestURI();
-		Map<String, String> labels = null;
-		if (StringUtils.isNotBlank(method)) {
-			labels = new HashMap<>();
-			labels.put("method", method);
+
+		Map<String, String> 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<String, String> 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 {

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<String, String> resolve(ServerWebExchange exchange);
+
+}

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<String, String> resolve(HttpServletRequest request);
+
+}

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<String, String> resolve(HttpServletRequest request) {
+		//rate limit by some request params. such as query params, headers ..
+
+		Map<String, String> labels = new HashMap<>();
+		labels.put("user", "zhangsan");
+
+		return labels;
+	}
+}