support parse ratelimit rule expression labels

CHANGELOG.md

@@ -1,4 +1,4 @@
 # Change Log
 ---
 
-- [fix:fix wrong isAliveFlag config when creating new PolarisServer.](https://github.com/Tencent/spring-cloud-tencent/pull/179)
+- [Feature: Support parse ratelimit rule expression labels.](https://github.com/Tencent/spring-cloud-tencent/pull/182)

changes/changes-1.4.4.md

@@ -0,0 +1,4 @@
+# Change Log
+---
+
+- [fix:fix wrong isAliveFlag config when creating new PolarisServer.](https://github.com/Tencent/spring-cloud-tencent/pull/179)

spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/RateLimitRuleLabelResolver.java

@@ -0,0 +1,72 @@
+/*
+ * 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;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.tencent.cloud.common.util.ExpressionLabelUtils;
+import com.tencent.cloud.polaris.context.ServiceRuleManager;
+import com.tencent.polaris.client.pb.ModelProto;
+import com.tencent.polaris.client.pb.RateLimitProto;
+
+import org.springframework.util.CollectionUtils;
+
+/**
+ * resolve labels from rate limit rule.
+ *
+ *@author lepdou 2022-05-13
+ */
+public class RateLimitRuleLabelResolver {
+
+	private final ServiceRuleManager serviceRuleManager;
+
+	public RateLimitRuleLabelResolver(ServiceRuleManager serviceRuleManager) {
+		this.serviceRuleManager = serviceRuleManager;
+	}
+
+	public Set<String> getExpressionLabelKeys(String namespace, String service) {
+		RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service);
+		if (rateLimitRule == null) {
+			return Collections.emptySet();
+		}
+
+		List<RateLimitProto.Rule> rules = rateLimitRule.getRulesList();
+		if (CollectionUtils.isEmpty(rules)) {
+			return Collections.emptySet();
+		}
+
+		Set<String> expressionLabels = new HashSet<>();
+		for (RateLimitProto.Rule rule : rules) {
+			Map<String, ModelProto.MatchString> labels = rule.getLabelsMap();
+			if (CollectionUtils.isEmpty(labels)) {
+				return Collections.emptySet();
+			}
+			for (String key : labels.keySet()) {
+				if (ExpressionLabelUtils.isExpressionLabel(key)) {
+					expressionLabels.add(key);
+				}
+			}
+		}
+		return expressionLabels;
+	}
+}

spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/config/RateLimitConfiguration.java

@@ -19,6 +19,8 @@
 package com.tencent.cloud.polaris.ratelimit.config;
 
 import com.tencent.cloud.polaris.context.ConditionalOnPolarisEnabled;
+import com.tencent.cloud.polaris.context.ServiceRuleManager;
+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;
@@ -62,6 +64,11 @@ public LimitAPI limitAPI(SDKContext polarisContext) {
 		return LimitAPIFactory.createLimitAPIByContext(polarisContext);
 	}
 
+	@Bean
+	public RateLimitRuleLabelResolver rateLimitRuleLabelService(ServiceRuleManager serviceRuleManager) {
+		return new RateLimitRuleLabelResolver(serviceRuleManager);
+	}
+
 	/**
 	 * Create when web application type is SERVLET.
 	 */
@@ -73,9 +80,10 @@ static class QuotaCheckFilterConfig {
 		@ConditionalOnMissingBean
 		public QuotaCheckServletFilter quotaCheckFilter(LimitAPI limitAPI,
 				@Nullable PolarisRateLimiterLabelServletResolver labelResolver,
-				PolarisRateLimitProperties polarisRateLimitProperties) {
+				PolarisRateLimitProperties polarisRateLimitProperties,
+				RateLimitRuleLabelResolver rateLimitRuleLabelResolver) {
 			return new QuotaCheckServletFilter(limitAPI, labelResolver,
-					polarisRateLimitProperties);
+					polarisRateLimitProperties, rateLimitRuleLabelResolver);
 		}
 
 		@Bean
@@ -101,9 +109,10 @@ static class MetadataReactiveFilterConfig {
 		@Bean
 		public QuotaCheckReactiveFilter quotaCheckReactiveFilter(LimitAPI limitAPI,
 				@Nullable PolarisRateLimiterLabelReactiveResolver labelResolver,
-				PolarisRateLimitProperties polarisRateLimitProperties) {
+				PolarisRateLimitProperties polarisRateLimitProperties,
+				RateLimitRuleLabelResolver rateLimitRuleLabelResolver) {
 			return new QuotaCheckReactiveFilter(limitAPI, labelResolver,
-					polarisRateLimitProperties);
+					polarisRateLimitProperties, rateLimitRuleLabelResolver);
 		}
 
 	}

spring-cloud-starter-tencent-polaris-ratelimit/src/main/java/com/tencent/cloud/polaris/ratelimit/filter/QuotaCheckReactiveFilter.java

@@ -21,15 +21,19 @@
 import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.Map;
+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.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.utils.QuotaCheckUtils;
 import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils;
+import com.tencent.cloud.common.util.ExpressionLabelUtils;
 import com.tencent.polaris.ratelimit.api.core.LimitAPI;
 import com.tencent.polaris.ratelimit.api.rpc.QuotaResponse;
 import com.tencent.polaris.ratelimit.api.rpc.QuotaResultCode;
@@ -42,7 +46,6 @@
 import org.springframework.core.io.buffer.DataBuffer;
 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;
@@ -52,7 +55,7 @@
 /**
  * Reactive filter to check quota.
  *
- * @author Haotian Zhang
+ * @author Haotian Zhang, lepdou
  */
 public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
 
@@ -65,14 +68,18 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
 
 	private final PolarisRateLimitProperties polarisRateLimitProperties;
 
+	private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver;
+
 	private String rejectTips;
 
 	public QuotaCheckReactiveFilter(LimitAPI limitAPI,
 			PolarisRateLimiterLabelReactiveResolver labelResolver,
-			PolarisRateLimitProperties polarisRateLimitProperties) {
+			PolarisRateLimitProperties polarisRateLimitProperties,
+			RateLimitRuleLabelResolver rateLimitRuleLabelResolver) {
 		this.limitAPI = limitAPI;
 		this.labelResolver = labelResolver;
 		this.polarisRateLimitProperties = polarisRateLimitProperties;
+		this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver;
 	}
 
 	@PostConstruct
@@ -90,27 +97,7 @@ public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
 		String localNamespace = MetadataContext.LOCAL_NAMESPACE;
 		String localService = MetadataContext.LOCAL_SERVICE;
 
-		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);
-			}
-		}
+		Map<String, String> labels = getRequestLabels(exchange, localNamespace, localService);
 
 		try {
 			QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI,
@@ -134,4 +121,42 @@ public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
 		return chain.filter(exchange);
 	}
 
+	private Map<String, String> getRequestLabels(ServerWebExchange exchange, String localNamespace, String localService) {
+		Map<String, String> labels = new HashMap<>();
+
+		// add build in labels
+		String path = exchange.getRequest().getURI().getPath();
+		if (StringUtils.isNotBlank(path)) {
+			labels.put(LABEL_METHOD, path);
+		}
+
+		// add rule expression labels
+		Map<String, String> expressionLabels = getRuleExpressionLabels(exchange, localNamespace, localService);
+		labels.putAll(expressionLabels);
+
+		// add custom labels
+		Map<String, String> customResolvedLabels = getCustomResolvedLabels(exchange);
+		labels.putAll(customResolvedLabels);
+
+		return labels;
+	}
+
+	private Map<String, String> getCustomResolvedLabels(ServerWebExchange exchange) {
+		if (labelResolver != null) {
+			try {
+				return labelResolver.resolve(exchange);
+			}
+			catch (Throwable e) {
+				LOG.error("resolve custom label failed. resolver = {}",
+						labelResolver.getClass().getName(), e);
+			}
+		}
+		return Maps.newHashMap();
+	}
+
+	private Map<String, String> getRuleExpressionLabels(ServerWebExchange exchange, String namespace, String service) {
+		Set<String> expressionLabels = rateLimitRuleLabelResolver.getExpressionLabelKeys(namespace, service);
+		return ExpressionLabelUtils.resolve(exchange, expressionLabels);
+	}
+
 }