feat:support concurrency rate limit. (#1455)
* feat:support concurrency rate limit. * feat:support concurrency rate limit.pull/1479/head
Haotian Zhang
1 month ago
committed by
GitHub
parent
7e9b96bae1
commit
ec43a2528c
@ -1,73 +0,0 @@
|
|||||||
/*
|
|
||||||
* 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.expresstion.ExpressionLabelUtils;
|
|
||||||
import com.tencent.cloud.polaris.context.ServiceRuleManager;
|
|
||||||
import com.tencent.polaris.specification.api.v1.model.ModelProto;
|
|
||||||
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
|
|
||||||
|
|
||||||
import org.springframework.util.CollectionUtils;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* resolve labels from rate limit rule.
|
|
||||||
*
|
|
||||||
*@author lepdou 2022-05-13
|
|
||||||
*/
|
|
||||||
@Deprecated
|
|
||||||
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;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@ -1,128 +0,0 @@
|
|||||||
/*
|
|
||||||
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
|
|
||||||
*
|
|
||||||
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
|
|
||||||
*
|
|
||||||
* Licensed under the BSD 3-Clause License (the "License");
|
|
||||||
* you may not use this file except in compliance with the License.
|
|
||||||
* You may obtain a copy of the License at
|
|
||||||
*
|
|
||||||
* https://opensource.org/licenses/BSD-3-Clause
|
|
||||||
*
|
|
||||||
* Unless required by applicable law or agreed to in writing, software distributed
|
|
||||||
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
|
|
||||||
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
|
||||||
* specific language governing permissions and limitations under the License.
|
|
||||||
*
|
|
||||||
*/
|
|
||||||
|
|
||||||
package com.tencent.cloud.polaris.ratelimit.resolver;
|
|
||||||
|
|
||||||
import java.net.InetSocketAddress;
|
|
||||||
import java.util.Collections;
|
|
||||||
import java.util.List;
|
|
||||||
import java.util.Map;
|
|
||||||
import java.util.Objects;
|
|
||||||
import java.util.Optional;
|
|
||||||
import java.util.Set;
|
|
||||||
import java.util.stream.Collectors;
|
|
||||||
|
|
||||||
import com.tencent.cloud.common.metadata.MetadataContextHolder;
|
|
||||||
import com.tencent.cloud.polaris.context.ServiceRuleManager;
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
|
|
||||||
import com.tencent.polaris.ratelimit.api.rpc.Argument;
|
|
||||||
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
|
|
||||||
import org.apache.commons.lang.StringUtils;
|
|
||||||
import org.slf4j.Logger;
|
|
||||||
import org.slf4j.LoggerFactory;
|
|
||||||
|
|
||||||
import org.springframework.util.CollectionUtils;
|
|
||||||
import org.springframework.web.server.ServerWebExchange;
|
|
||||||
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* resolve arguments from rate limit rule for Reactive.
|
|
||||||
*
|
|
||||||
* @author seansyyu 2023-03-09
|
|
||||||
*/
|
|
||||||
public class RateLimitRuleArgumentReactiveResolver {
|
|
||||||
|
|
||||||
private static final Logger LOG = LoggerFactory.getLogger(RateLimitRuleArgumentReactiveResolver.class);
|
|
||||||
|
|
||||||
private final ServiceRuleManager serviceRuleManager;
|
|
||||||
|
|
||||||
private final PolarisRateLimiterLabelReactiveResolver labelResolver;
|
|
||||||
|
|
||||||
public RateLimitRuleArgumentReactiveResolver(ServiceRuleManager serviceRuleManager, PolarisRateLimiterLabelReactiveResolver labelResolver) {
|
|
||||||
this.serviceRuleManager = serviceRuleManager;
|
|
||||||
this.labelResolver = labelResolver;
|
|
||||||
}
|
|
||||||
|
|
||||||
public Set<Argument> getArguments(ServerWebExchange request, String namespace, String service) {
|
|
||||||
RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service);
|
|
||||||
if (rateLimitRule == null) {
|
|
||||||
return Collections.emptySet();
|
|
||||||
}
|
|
||||||
List<RateLimitProto.Rule> rules = rateLimitRule.getRulesList();
|
|
||||||
if (CollectionUtils.isEmpty(rules)) {
|
|
||||||
return Collections.emptySet();
|
|
||||||
}
|
|
||||||
return rules.stream()
|
|
||||||
.flatMap(rule -> rule.getArgumentsList().stream())
|
|
||||||
.map(matchArgument -> {
|
|
||||||
String matchKey = matchArgument.getKey();
|
|
||||||
Argument argument = null;
|
|
||||||
switch (matchArgument.getType()) {
|
|
||||||
case CUSTOM:
|
|
||||||
argument = StringUtils.isBlank(matchKey) ? null :
|
|
||||||
Argument.buildCustom(matchKey, Optional.ofNullable(getCustomResolvedLabels(request).get(matchKey))
|
|
||||||
.orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
case METHOD:
|
|
||||||
argument = Argument.buildMethod(request.getRequest().getMethod().name());
|
|
||||||
break;
|
|
||||||
case HEADER:
|
|
||||||
argument = StringUtils.isBlank(matchKey) ? null :
|
|
||||||
Argument.buildHeader(matchKey, Optional.ofNullable(request.getRequest().getHeaders()
|
|
||||||
.getFirst(matchKey)).orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
case QUERY:
|
|
||||||
argument = StringUtils.isBlank(matchKey) ? null :
|
|
||||||
Argument.buildQuery(matchKey, Optional.ofNullable(request.getRequest().getQueryParams()
|
|
||||||
.getFirst(matchKey)).orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
case CALLER_SERVICE:
|
|
||||||
String sourceServiceNamespace = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, true)
|
|
||||||
.orElse(StringUtils.EMPTY);
|
|
||||||
String sourceServiceName = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAME, true)
|
|
||||||
.orElse(StringUtils.EMPTY);
|
|
||||||
if (!StringUtils.isEmpty(sourceServiceNamespace) && !StringUtils.isEmpty(sourceServiceName)) {
|
|
||||||
argument = Argument.buildCallerService(sourceServiceNamespace, sourceServiceName);
|
|
||||||
}
|
|
||||||
break;
|
|
||||||
case CALLER_IP:
|
|
||||||
InetSocketAddress remoteAddress = request.getRequest().getRemoteAddress();
|
|
||||||
argument = Argument.buildCallerIP(remoteAddress != null ? remoteAddress.getAddress()
|
|
||||||
.getHostAddress() : StringUtils.EMPTY);
|
|
||||||
break;
|
|
||||||
default:
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
return argument;
|
|
||||||
}).filter(Objects::nonNull).collect(Collectors.toSet());
|
|
||||||
}
|
|
||||||
|
|
||||||
private Map<String, String> getCustomResolvedLabels(ServerWebExchange request) {
|
|
||||||
if (labelResolver != null) {
|
|
||||||
try {
|
|
||||||
return labelResolver.resolve(request);
|
|
||||||
}
|
|
||||||
catch (Throwable e) {
|
|
||||||
LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return Collections.emptyMap();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@ -1,122 +0,0 @@
|
|||||||
/*
|
|
||||||
* 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 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 jakarta.servlet.http.HttpServletRequest;
|
|
||||||
import org.apache.commons.lang.StringUtils;
|
|
||||||
import org.slf4j.Logger;
|
|
||||||
import org.slf4j.LoggerFactory;
|
|
||||||
|
|
||||||
import org.springframework.util.CollectionUtils;
|
|
||||||
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* resolve arguments from rate limit rule for Servlet.
|
|
||||||
*
|
|
||||||
* @author seansyyu 2023-03-09
|
|
||||||
*/
|
|
||||||
public class RateLimitRuleArgumentServletResolver {
|
|
||||||
|
|
||||||
private static final Logger LOG = LoggerFactory.getLogger(QuotaCheckServletFilter.class);
|
|
||||||
|
|
||||||
private final ServiceRuleManager serviceRuleManager;
|
|
||||||
|
|
||||||
private final PolarisRateLimiterLabelServletResolver labelResolver;
|
|
||||||
|
|
||||||
public RateLimitRuleArgumentServletResolver(ServiceRuleManager serviceRuleManager, PolarisRateLimiterLabelServletResolver labelResolver) {
|
|
||||||
this.serviceRuleManager = serviceRuleManager;
|
|
||||||
this.labelResolver = labelResolver;
|
|
||||||
}
|
|
||||||
|
|
||||||
public Set<Argument> getArguments(HttpServletRequest request, String namespace, String service) {
|
|
||||||
RateLimitProto.RateLimit rateLimitRule = serviceRuleManager.getServiceRateLimitRule(namespace, service);
|
|
||||||
if (rateLimitRule == null) {
|
|
||||||
return Collections.emptySet();
|
|
||||||
}
|
|
||||||
List<RateLimitProto.Rule> rules = rateLimitRule.getRulesList();
|
|
||||||
if (CollectionUtils.isEmpty(rules)) {
|
|
||||||
return Collections.emptySet();
|
|
||||||
}
|
|
||||||
return rules.stream()
|
|
||||||
.flatMap(rule -> rule.getArgumentsList().stream())
|
|
||||||
.map(matchArgument -> {
|
|
||||||
String matchKey = matchArgument.getKey();
|
|
||||||
Argument argument = null;
|
|
||||||
switch (matchArgument.getType()) {
|
|
||||||
case CUSTOM:
|
|
||||||
argument = StringUtils.isBlank(matchKey) ? null :
|
|
||||||
Argument.buildCustom(matchKey, Optional.ofNullable(getCustomResolvedLabels(request).get(matchKey)).orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
case METHOD:
|
|
||||||
argument = Argument.buildMethod(request.getMethod());
|
|
||||||
break;
|
|
||||||
case HEADER:
|
|
||||||
argument = StringUtils.isBlank(matchKey) ? null :
|
|
||||||
Argument.buildHeader(matchKey, Optional.ofNullable(request.getHeader(matchKey)).orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
case QUERY:
|
|
||||||
argument = StringUtils.isBlank(matchKey) ? null :
|
|
||||||
Argument.buildQuery(matchKey, Optional.ofNullable(request.getParameter(matchKey)).orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
case CALLER_SERVICE:
|
|
||||||
String sourceServiceNamespace = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, true).orElse(StringUtils.EMPTY);
|
|
||||||
String sourceServiceName = MetadataContextHolder.getDisposableMetadata(DEFAULT_METADATA_SOURCE_SERVICE_NAME, true).orElse(StringUtils.EMPTY);
|
|
||||||
if (!StringUtils.isEmpty(sourceServiceNamespace) && !StringUtils.isEmpty(sourceServiceName)) {
|
|
||||||
argument = Argument.buildCallerService(sourceServiceNamespace, sourceServiceName);
|
|
||||||
}
|
|
||||||
break;
|
|
||||||
case CALLER_IP:
|
|
||||||
argument = Argument.buildCallerIP(Optional.ofNullable(request.getRemoteAddr()).orElse(StringUtils.EMPTY));
|
|
||||||
break;
|
|
||||||
default:
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
return argument;
|
|
||||||
}).filter(Objects::nonNull).collect(Collectors.toSet());
|
|
||||||
}
|
|
||||||
|
|
||||||
private Map<String, String> getCustomResolvedLabels(HttpServletRequest request) {
|
|
||||||
if (labelResolver != null) {
|
|
||||||
try {
|
|
||||||
return labelResolver.resolve(request);
|
|
||||||
}
|
|
||||||
catch (Throwable e) {
|
|
||||||
LOG.error("resolve custom label failed. resolver = {}", labelResolver.getClass().getName(), e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return Collections.emptyMap();
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
@ -1,38 +0,0 @@
|
|||||||
/*
|
|
||||||
* 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);
|
|
||||||
}
|
|
||||||
@ -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.polaris.ratelimit.tsf;
|
||||||
|
|
||||||
|
import com.tencent.cloud.common.tsf.ConditionalOnTsfConsulEnabled;
|
||||||
|
import com.tencent.cloud.polaris.context.config.extend.consul.ConsulProperties;
|
||||||
|
import com.tencent.cloud.polaris.context.config.extend.tsf.TsfCoreProperties;
|
||||||
|
|
||||||
|
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
|
||||||
|
import org.springframework.context.annotation.Bean;
|
||||||
|
import org.springframework.context.annotation.Configuration;
|
||||||
|
import org.springframework.core.env.Environment;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Auto configuration of TSF rate limit.
|
||||||
|
*
|
||||||
|
* @author Haotian Zhang
|
||||||
|
*/
|
||||||
|
@Configuration(proxyBeanMethods = false)
|
||||||
|
@ConditionalOnTsfConsulEnabled
|
||||||
|
public class TsfRateLimitAutoConfiguration {
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
@ConditionalOnMissingBean
|
||||||
|
public TsfRateLimitConfigModifier tsfRateLimitConfigModifier(TsfCoreProperties tsfCoreProperties,
|
||||||
|
ConsulProperties consulProperties, Environment environment) {
|
||||||
|
return new TsfRateLimitConfigModifier(tsfCoreProperties, consulProperties, environment);
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,68 @@
|
|||||||
|
/*
|
||||||
|
* 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.tsf;
|
||||||
|
|
||||||
|
import java.util.Map;
|
||||||
|
|
||||||
|
import com.tencent.cloud.common.constant.OrderConstant;
|
||||||
|
import com.tencent.cloud.polaris.context.PolarisConfigModifier;
|
||||||
|
import com.tencent.cloud.polaris.context.config.extend.consul.ConsulProperties;
|
||||||
|
import com.tencent.cloud.polaris.context.config.extend.tsf.TsfContextUtils;
|
||||||
|
import com.tencent.cloud.polaris.context.config.extend.tsf.TsfCoreProperties;
|
||||||
|
import com.tencent.polaris.factory.config.ConfigurationImpl;
|
||||||
|
import com.tencent.polaris.ratelimit.client.sync.tsf.TsfRateLimitConstants;
|
||||||
|
|
||||||
|
import org.springframework.core.env.Environment;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Config modifier for TSF rate limit.
|
||||||
|
*
|
||||||
|
* @author Haotian Zhang
|
||||||
|
*/
|
||||||
|
public class TsfRateLimitConfigModifier implements PolarisConfigModifier {
|
||||||
|
|
||||||
|
private final TsfCoreProperties tsfCoreProperties;
|
||||||
|
|
||||||
|
private final ConsulProperties consulProperties;
|
||||||
|
|
||||||
|
private final Environment environment;
|
||||||
|
|
||||||
|
public TsfRateLimitConfigModifier(TsfCoreProperties tsfCoreProperties, ConsulProperties consulProperties,
|
||||||
|
Environment environment) {
|
||||||
|
this.tsfCoreProperties = tsfCoreProperties;
|
||||||
|
this.consulProperties = consulProperties;
|
||||||
|
this.environment = environment;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void modify(ConfigurationImpl configuration) {
|
||||||
|
if (TsfContextUtils.isTsfConsulEnabled(environment)) {
|
||||||
|
Map<String, String> metadata = configuration.getProvider().getRateLimit().getMetadata();
|
||||||
|
metadata.put(TsfRateLimitConstants.RATE_LIMIT_MASTER_IP_KEY, tsfCoreProperties.getRatelimitMasterIp());
|
||||||
|
metadata.put(TsfRateLimitConstants.RATE_LIMIT_MASTER_PORT_KEY, String.valueOf(tsfCoreProperties.getRatelimitMasterPort()));
|
||||||
|
metadata.put(TsfRateLimitConstants.SERVICE_NAME_KEY, tsfCoreProperties.getServiceName());
|
||||||
|
metadata.put(TsfRateLimitConstants.INSTANCE_ID_KEY, tsfCoreProperties.getInstanceId());
|
||||||
|
metadata.put(TsfRateLimitConstants.TOKEN_KEY, consulProperties.getAclToken());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return OrderConstant.Modifier.RATE_LIMIT_ORDER;
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -1,2 +1,3 @@
|
|||||||
org.springframework.cloud.bootstrap.BootstrapConfiguration=\
|
org.springframework.cloud.bootstrap.BootstrapConfiguration=\
|
||||||
com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitPropertiesBootstrapConfiguration
|
com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitPropertiesBootstrapConfiguration,\
|
||||||
|
com.tencent.cloud.polaris.ratelimit.tsf.TsfRateLimitBootstrapConfiguration
|
||||||
|
|||||||
@ -1,3 +1,4 @@
|
|||||||
com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitAutoConfiguration
|
com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitAutoConfiguration
|
||||||
com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitPropertiesAutoConfiguration
|
com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitPropertiesAutoConfiguration
|
||||||
com.tencent.cloud.polaris.ratelimit.endpoint.PolarisRateLimitRuleEndpointAutoConfiguration
|
com.tencent.cloud.polaris.ratelimit.endpoint.PolarisRateLimitRuleEndpointAutoConfiguration
|
||||||
|
com.tencent.cloud.polaris.ratelimit.tsf.TsfRateLimitAutoConfiguration
|
||||||
|
|||||||
@ -1,99 +0,0 @@
|
|||||||
/*
|
|
||||||
* 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.Set;
|
|
||||||
|
|
||||||
import com.google.protobuf.StringValue;
|
|
||||||
import com.tencent.cloud.polaris.context.ServiceRuleManager;
|
|
||||||
import com.tencent.polaris.specification.api.v1.model.ModelProto;
|
|
||||||
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.junit.jupiter.MockitoExtension;
|
|
||||||
|
|
||||||
import static org.assertj.core.api.Assertions.assertThat;
|
|
||||||
import static org.mockito.ArgumentMatchers.any;
|
|
||||||
import static org.mockito.ArgumentMatchers.anyString;
|
|
||||||
import static org.mockito.Mockito.mock;
|
|
||||||
import static org.mockito.Mockito.when;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Test for {@link RateLimitRuleLabelResolver}.
|
|
||||||
*
|
|
||||||
* @author Haotian Zhang
|
|
||||||
*/
|
|
||||||
@ExtendWith(MockitoExtension.class)
|
|
||||||
public class RateLimitRuleLabelResolverTest {
|
|
||||||
|
|
||||||
private RateLimitRuleLabelResolver rateLimitRuleLabelResolver;
|
|
||||||
|
|
||||||
@BeforeEach
|
|
||||||
void setUp() {
|
|
||||||
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
|
|
||||||
when(serviceRuleManager.getServiceRateLimitRule(any(), anyString())).thenAnswer(invocationOnMock -> {
|
|
||||||
String serviceName = invocationOnMock.getArgument(1).toString();
|
|
||||||
if (serviceName.equals("TestApp1")) {
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
else if (serviceName.equals("TestApp2")) {
|
|
||||||
return RateLimitProto.RateLimit.newBuilder().build();
|
|
||||||
}
|
|
||||||
else if (serviceName.equals("TestApp3")) {
|
|
||||||
RateLimitProto.Rule rule = RateLimitProto.Rule.newBuilder().build();
|
|
||||||
return RateLimitProto.RateLimit.newBuilder().addRules(rule).build();
|
|
||||||
}
|
|
||||||
else {
|
|
||||||
ModelProto.MatchString matchString = ModelProto.MatchString.newBuilder()
|
|
||||||
.setType(ModelProto.MatchString.MatchStringType.EXACT)
|
|
||||||
.setValue(StringValue.of("value"))
|
|
||||||
.setValueType(ModelProto.MatchString.ValueType.TEXT).build();
|
|
||||||
RateLimitProto.Rule rule = RateLimitProto.Rule.newBuilder()
|
|
||||||
.putLabels("${http.method}", matchString).build();
|
|
||||||
return RateLimitProto.RateLimit.newBuilder().addRules(rule).build();
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
rateLimitRuleLabelResolver = new RateLimitRuleLabelResolver(serviceRuleManager);
|
|
||||||
}
|
|
||||||
|
|
||||||
@Test
|
|
||||||
public void testGetExpressionLabelKeys() {
|
|
||||||
// rateLimitRule == null
|
|
||||||
String serviceName = "TestApp1";
|
|
||||||
Set<String> labelKeys = rateLimitRuleLabelResolver.getExpressionLabelKeys(null, serviceName);
|
|
||||||
assertThat(labelKeys).isEmpty();
|
|
||||||
|
|
||||||
// CollectionUtils.isEmpty(rules)
|
|
||||||
serviceName = "TestApp2";
|
|
||||||
labelKeys = rateLimitRuleLabelResolver.getExpressionLabelKeys(null, serviceName);
|
|
||||||
assertThat(labelKeys).isEmpty();
|
|
||||||
|
|
||||||
// CollectionUtils.isEmpty(labels)
|
|
||||||
serviceName = "TestApp3";
|
|
||||||
labelKeys = rateLimitRuleLabelResolver.getExpressionLabelKeys(null, serviceName);
|
|
||||||
assertThat(labelKeys).isEmpty();
|
|
||||||
|
|
||||||
// Has labels
|
|
||||||
serviceName = "TestApp4";
|
|
||||||
labelKeys = rateLimitRuleLabelResolver.getExpressionLabelKeys(null, serviceName);
|
|
||||||
assertThat(labelKeys).isNotEmpty();
|
|
||||||
assertThat(labelKeys).contains("${http.method}");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@ -1,144 +0,0 @@
|
|||||||
/*
|
|
||||||
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
|
|
||||||
*
|
|
||||||
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
|
|
||||||
*
|
|
||||||
* Licensed under the BSD 3-Clause License (the "License");
|
|
||||||
* you may not use this file except in compliance with the License.
|
|
||||||
* You may obtain a copy of the License at
|
|
||||||
*
|
|
||||||
* https://opensource.org/licenses/BSD-3-Clause
|
|
||||||
*
|
|
||||||
* Unless required by applicable law or agreed to in writing, software distributed
|
|
||||||
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
|
|
||||||
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
|
||||||
* specific language governing permissions and limitations under the License.
|
|
||||||
*/
|
|
||||||
|
|
||||||
package com.tencent.cloud.polaris.ratelimit.resolver;
|
|
||||||
|
|
||||||
import java.io.BufferedReader;
|
|
||||||
import java.io.InputStream;
|
|
||||||
import java.io.InputStreamReader;
|
|
||||||
import java.net.InetSocketAddress;
|
|
||||||
import java.nio.charset.StandardCharsets;
|
|
||||||
import java.util.Collections;
|
|
||||||
import java.util.HashMap;
|
|
||||||
import java.util.HashSet;
|
|
||||||
import java.util.Set;
|
|
||||||
import java.util.stream.Collectors;
|
|
||||||
|
|
||||||
import com.google.protobuf.InvalidProtocolBufferException;
|
|
||||||
import com.google.protobuf.util.JsonFormat;
|
|
||||||
import com.tencent.cloud.common.metadata.MetadataContext;
|
|
||||||
import com.tencent.cloud.common.metadata.MetadataContextHolder;
|
|
||||||
import com.tencent.cloud.polaris.context.ServiceRuleManager;
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilterTest;
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
|
|
||||||
import com.tencent.polaris.ratelimit.api.rpc.Argument;
|
|
||||||
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
|
|
||||||
import org.junit.jupiter.api.BeforeEach;
|
|
||||||
import org.junit.jupiter.api.Test;
|
|
||||||
import org.junit.jupiter.api.extension.ExtendWith;
|
|
||||||
|
|
||||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
|
||||||
import org.springframework.boot.test.context.SpringBootTest;
|
|
||||||
import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
|
|
||||||
import org.springframework.mock.web.server.MockServerWebExchange;
|
|
||||||
import org.springframework.test.context.junit.jupiter.SpringExtension;
|
|
||||||
import org.springframework.web.server.ServerWebExchange;
|
|
||||||
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
|
|
||||||
import static org.assertj.core.api.Assertions.assertThat;
|
|
||||||
import static org.mockito.ArgumentMatchers.anyString;
|
|
||||||
import static org.mockito.Mockito.mock;
|
|
||||||
import static org.mockito.Mockito.when;
|
|
||||||
|
|
||||||
@ExtendWith(SpringExtension.class)
|
|
||||||
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
|
|
||||||
classes = RateLimitRuleArgumentReactiveResolverTest.TestApplication.class,
|
|
||||||
properties = {
|
|
||||||
"spring.cloud.polaris.namespace=Test", "spring.cloud.polaris.service=TestApp"
|
|
||||||
})
|
|
||||||
public class RateLimitRuleArgumentReactiveResolverTest {
|
|
||||||
|
|
||||||
private final PolarisRateLimiterLabelReactiveResolver labelResolver =
|
|
||||||
exchange -> Collections.singletonMap("xxx", "xxx");
|
|
||||||
|
|
||||||
private final PolarisRateLimiterLabelReactiveResolver labelResolverEx =
|
|
||||||
exchange -> {
|
|
||||||
throw new RuntimeException();
|
|
||||||
};
|
|
||||||
|
|
||||||
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver1;
|
|
||||||
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver2;
|
|
||||||
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver3;
|
|
||||||
private RateLimitRuleArgumentReactiveResolver rateLimitRuleArgumentReactiveResolver4;
|
|
||||||
|
|
||||||
@BeforeEach
|
|
||||||
void setUp() throws InvalidProtocolBufferException {
|
|
||||||
MetadataContext.LOCAL_NAMESPACE = "TEST";
|
|
||||||
|
|
||||||
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
|
|
||||||
|
|
||||||
RateLimitProto.Rule.Builder ratelimitRuleBuilder = RateLimitProto.Rule.newBuilder();
|
|
||||||
InputStream inputStream = QuotaCheckServletFilterTest.class.getClassLoader().getResourceAsStream("ratelimit.json");
|
|
||||||
String json = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)).lines().collect(Collectors.joining(""));
|
|
||||||
JsonFormat.parser().ignoringUnknownFields().merge(json, ratelimitRuleBuilder);
|
|
||||||
RateLimitProto.Rule rateLimitRule = ratelimitRuleBuilder.build();
|
|
||||||
RateLimitProto.RateLimit rateLimit = RateLimitProto.RateLimit.newBuilder().addRules(rateLimitRule).build();
|
|
||||||
when(serviceRuleManager.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit);
|
|
||||||
|
|
||||||
// normal
|
|
||||||
this.rateLimitRuleArgumentReactiveResolver1 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager, labelResolver);
|
|
||||||
// ex
|
|
||||||
this.rateLimitRuleArgumentReactiveResolver2 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager, labelResolverEx);
|
|
||||||
// null
|
|
||||||
ServiceRuleManager serviceRuleManager1 = mock(ServiceRuleManager.class);
|
|
||||||
when(serviceRuleManager1.getServiceRateLimitRule(anyString(), anyString())).thenReturn(null);
|
|
||||||
this.rateLimitRuleArgumentReactiveResolver3 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager1, labelResolver);
|
|
||||||
// null 2
|
|
||||||
ServiceRuleManager serviceRuleManager2 = mock(ServiceRuleManager.class);
|
|
||||||
RateLimitProto.RateLimit rateLimit2 = RateLimitProto.RateLimit.newBuilder().build();
|
|
||||||
when(serviceRuleManager2.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit2);
|
|
||||||
this.rateLimitRuleArgumentReactiveResolver4 = new RateLimitRuleArgumentReactiveResolver(serviceRuleManager2, labelResolver);
|
|
||||||
}
|
|
||||||
|
|
||||||
@Test
|
|
||||||
public void testGetRuleArguments() {
|
|
||||||
// Mock request
|
|
||||||
MetadataContext.LOCAL_SERVICE = "Test";
|
|
||||||
// Mock request
|
|
||||||
MockServerHttpRequest request = MockServerHttpRequest.get("http://127.0.0.1:8080/test")
|
|
||||||
.remoteAddress(new InetSocketAddress("127.0.0.1", 8080))
|
|
||||||
.header("xxx", "xxx")
|
|
||||||
.queryParam("yyy", "yyy")
|
|
||||||
.build();
|
|
||||||
ServerWebExchange exchange = MockServerWebExchange.from(request);
|
|
||||||
MetadataContext metadataContext = new MetadataContext();
|
|
||||||
metadataContext.setUpstreamDisposableMetadata(new HashMap<String, String>() {{
|
|
||||||
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE);
|
|
||||||
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE);
|
|
||||||
}});
|
|
||||||
MetadataContextHolder.set(metadataContext);
|
|
||||||
Set<Argument> arguments = rateLimitRuleArgumentReactiveResolver1.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
Set<Argument> exceptRes = new HashSet<>();
|
|
||||||
exceptRes.add(Argument.buildMethod("GET"));
|
|
||||||
exceptRes.add(Argument.buildHeader("xxx", "xxx"));
|
|
||||||
exceptRes.add(Argument.buildQuery("yyy", "yyy"));
|
|
||||||
exceptRes.add(Argument.buildCallerIP("127.0.0.1"));
|
|
||||||
exceptRes.add(Argument.buildCustom("xxx", "xxx"));
|
|
||||||
exceptRes.add(Argument.buildCallerService(MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE));
|
|
||||||
assertThat(arguments).isEqualTo(exceptRes);
|
|
||||||
|
|
||||||
rateLimitRuleArgumentReactiveResolver2.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
rateLimitRuleArgumentReactiveResolver3.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
rateLimitRuleArgumentReactiveResolver4.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
}
|
|
||||||
|
|
||||||
@SpringBootApplication
|
|
||||||
protected static class TestApplication {
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
@ -1,135 +0,0 @@
|
|||||||
/*
|
|
||||||
* Tencent is pleased to support the open source community by making Spring Cloud Tencent available.
|
|
||||||
*
|
|
||||||
* Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
|
|
||||||
*
|
|
||||||
* Licensed under the BSD 3-Clause License (the "License");
|
|
||||||
* you may not use this file except in compliance with the License.
|
|
||||||
* You may obtain a copy of the License at
|
|
||||||
*
|
|
||||||
* https://opensource.org/licenses/BSD-3-Clause
|
|
||||||
*
|
|
||||||
* Unless required by applicable law or agreed to in writing, software distributed
|
|
||||||
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
|
|
||||||
* CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
|
||||||
* specific language governing permissions and limitations under the License.
|
|
||||||
*/
|
|
||||||
|
|
||||||
package com.tencent.cloud.polaris.ratelimit.resolver;
|
|
||||||
|
|
||||||
import java.io.BufferedReader;
|
|
||||||
import java.io.InputStream;
|
|
||||||
import java.io.InputStreamReader;
|
|
||||||
import java.nio.charset.StandardCharsets;
|
|
||||||
import java.util.Collections;
|
|
||||||
import java.util.HashMap;
|
|
||||||
import java.util.HashSet;
|
|
||||||
import java.util.Set;
|
|
||||||
import java.util.stream.Collectors;
|
|
||||||
|
|
||||||
import com.google.protobuf.InvalidProtocolBufferException;
|
|
||||||
import com.google.protobuf.util.JsonFormat;
|
|
||||||
import com.tencent.cloud.common.metadata.MetadataContext;
|
|
||||||
import com.tencent.cloud.common.metadata.MetadataContextHolder;
|
|
||||||
import com.tencent.cloud.polaris.context.ServiceRuleManager;
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilterTest;
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
|
|
||||||
import com.tencent.polaris.ratelimit.api.rpc.Argument;
|
|
||||||
import com.tencent.polaris.specification.api.v1.traffic.manage.RateLimitProto;
|
|
||||||
import org.junit.jupiter.api.BeforeEach;
|
|
||||||
import org.junit.jupiter.api.Test;
|
|
||||||
import org.junit.jupiter.api.extension.ExtendWith;
|
|
||||||
|
|
||||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
|
||||||
import org.springframework.boot.test.context.SpringBootTest;
|
|
||||||
import org.springframework.mock.web.MockHttpServletRequest;
|
|
||||||
import org.springframework.test.context.junit.jupiter.SpringExtension;
|
|
||||||
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
|
|
||||||
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
|
|
||||||
import static org.assertj.core.api.Assertions.assertThat;
|
|
||||||
import static org.mockito.ArgumentMatchers.anyString;
|
|
||||||
import static org.mockito.Mockito.mock;
|
|
||||||
import static org.mockito.Mockito.when;
|
|
||||||
|
|
||||||
@ExtendWith(SpringExtension.class)
|
|
||||||
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
|
|
||||||
classes = RateLimitRuleArgumentServletResolverTest.TestApplication.class,
|
|
||||||
properties = {
|
|
||||||
"spring.cloud.polaris.namespace=Test", "spring.cloud.polaris.service=TestApp"
|
|
||||||
})
|
|
||||||
public class RateLimitRuleArgumentServletResolverTest {
|
|
||||||
|
|
||||||
private final PolarisRateLimiterLabelServletResolver labelResolver =
|
|
||||||
exchange -> Collections.singletonMap("xxx", "xxx");
|
|
||||||
private final PolarisRateLimiterLabelServletResolver labelResolverEx =
|
|
||||||
exchange -> {
|
|
||||||
throw new RuntimeException();
|
|
||||||
};
|
|
||||||
|
|
||||||
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver1;
|
|
||||||
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver2;
|
|
||||||
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver3;
|
|
||||||
private RateLimitRuleArgumentServletResolver rateLimitRuleArgumentServletResolver4;
|
|
||||||
|
|
||||||
@BeforeEach
|
|
||||||
void setUp() throws InvalidProtocolBufferException {
|
|
||||||
MetadataContext.LOCAL_NAMESPACE = "TEST";
|
|
||||||
|
|
||||||
ServiceRuleManager serviceRuleManager = mock(ServiceRuleManager.class);
|
|
||||||
|
|
||||||
RateLimitProto.Rule.Builder ratelimitRuleBuilder = RateLimitProto.Rule.newBuilder();
|
|
||||||
InputStream inputStream = QuotaCheckServletFilterTest.class.getClassLoader().getResourceAsStream("ratelimit.json");
|
|
||||||
String json = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)).lines().collect(Collectors.joining(""));
|
|
||||||
JsonFormat.parser().ignoringUnknownFields().merge(json, ratelimitRuleBuilder);
|
|
||||||
RateLimitProto.Rule rateLimitRule = ratelimitRuleBuilder.build();
|
|
||||||
RateLimitProto.RateLimit rateLimit = RateLimitProto.RateLimit.newBuilder().addRules(rateLimitRule).build();
|
|
||||||
when(serviceRuleManager.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit);
|
|
||||||
|
|
||||||
// normal
|
|
||||||
this.rateLimitRuleArgumentServletResolver1 = new RateLimitRuleArgumentServletResolver(serviceRuleManager, labelResolver);
|
|
||||||
// ex
|
|
||||||
this.rateLimitRuleArgumentServletResolver2 = new RateLimitRuleArgumentServletResolver(serviceRuleManager, labelResolverEx);
|
|
||||||
// null
|
|
||||||
ServiceRuleManager serviceRuleManager1 = mock(ServiceRuleManager.class);
|
|
||||||
when(serviceRuleManager1.getServiceRateLimitRule(anyString(), anyString())).thenReturn(null);
|
|
||||||
this.rateLimitRuleArgumentServletResolver3 = new RateLimitRuleArgumentServletResolver(serviceRuleManager1, labelResolver);
|
|
||||||
// null 2
|
|
||||||
ServiceRuleManager serviceRuleManager2 = mock(ServiceRuleManager.class);
|
|
||||||
RateLimitProto.RateLimit rateLimit2 = RateLimitProto.RateLimit.newBuilder().build();
|
|
||||||
when(serviceRuleManager2.getServiceRateLimitRule(anyString(), anyString())).thenReturn(rateLimit2);
|
|
||||||
this.rateLimitRuleArgumentServletResolver4 = new RateLimitRuleArgumentServletResolver(serviceRuleManager2, labelResolver);
|
|
||||||
}
|
|
||||||
|
|
||||||
@Test
|
|
||||||
public void testGetRuleArguments() {
|
|
||||||
// Mock request
|
|
||||||
MetadataContext.LOCAL_SERVICE = "Test";
|
|
||||||
MockHttpServletRequest request = new MockHttpServletRequest(null, "GET", "/xxx");
|
|
||||||
request.setParameter("yyy", "yyy");
|
|
||||||
request.addHeader("xxx", "xxx");
|
|
||||||
MetadataContext metadataContext = new MetadataContext();
|
|
||||||
metadataContext.setUpstreamDisposableMetadata(new HashMap<String, String>() {{
|
|
||||||
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE);
|
|
||||||
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE);
|
|
||||||
}});
|
|
||||||
MetadataContextHolder.set(metadataContext);
|
|
||||||
Set<Argument> arguments = rateLimitRuleArgumentServletResolver1.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
Set<Argument> exceptRes = new HashSet<>();
|
|
||||||
exceptRes.add(Argument.buildMethod("GET"));
|
|
||||||
exceptRes.add(Argument.buildHeader("xxx", "xxx"));
|
|
||||||
exceptRes.add(Argument.buildQuery("yyy", "yyy"));
|
|
||||||
exceptRes.add(Argument.buildCallerIP("127.0.0.1"));
|
|
||||||
exceptRes.add(Argument.buildCustom("xxx", "xxx"));
|
|
||||||
exceptRes.add(Argument.buildCallerService(MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE));
|
|
||||||
assertThat(arguments).isEqualTo(exceptRes);
|
|
||||||
|
|
||||||
rateLimitRuleArgumentServletResolver2.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
rateLimitRuleArgumentServletResolver3.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
rateLimitRuleArgumentServletResolver4.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
|
|
||||||
}
|
|
||||||
|
|
||||||
@SpringBootApplication
|
|
||||||
protected static class TestApplication {
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@ -1,62 +0,0 @@
|
|||||||
/*
|
|
||||||
* 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.quickstart.callee.ratelimit;
|
|
||||||
|
|
||||||
import java.util.HashMap;
|
|
||||||
import java.util.Map;
|
|
||||||
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
|
|
||||||
import jakarta.servlet.http.HttpServletRequest;
|
|
||||||
import org.slf4j.Logger;
|
|
||||||
import org.slf4j.LoggerFactory;
|
|
||||||
|
|
||||||
import org.springframework.beans.factory.annotation.Value;
|
|
||||||
import org.springframework.stereotype.Component;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* resolver custom label from request.
|
|
||||||
*
|
|
||||||
* @author lepdou 2022-03-31
|
|
||||||
*/
|
|
||||||
@Component
|
|
||||||
public class CustomLabelResolver implements PolarisRateLimiterLabelServletResolver {
|
|
||||||
private static final Logger LOG = LoggerFactory.getLogger(CustomLabelResolver.class);
|
|
||||||
|
|
||||||
@Value("${label.key-value:}")
|
|
||||||
private String[] keyValues;
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public Map<String, String> resolve(HttpServletRequest request) {
|
|
||||||
// rate limit by some request params. such as query params, headers ..
|
|
||||||
|
|
||||||
return getLabels(keyValues);
|
|
||||||
}
|
|
||||||
|
|
||||||
private Map<String, String> getLabels(String[] keyValues) {
|
|
||||||
Map<String, String> labels = new HashMap<>();
|
|
||||||
for (String kv : keyValues) {
|
|
||||||
String key = kv.substring(0, kv.indexOf(":"));
|
|
||||||
String value = kv.substring(kv.indexOf(":") + 1);
|
|
||||||
labels.put(key, value);
|
|
||||||
}
|
|
||||||
|
|
||||||
LOG.info("Current labels:{}", labels);
|
|
||||||
return labels;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@ -1,60 +0,0 @@
|
|||||||
/*
|
|
||||||
* 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.quickstart.callee.ratelimit;
|
|
||||||
|
|
||||||
import java.util.HashMap;
|
|
||||||
import java.util.Map;
|
|
||||||
|
|
||||||
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
|
|
||||||
import org.slf4j.Logger;
|
|
||||||
import org.slf4j.LoggerFactory;
|
|
||||||
|
|
||||||
import org.springframework.beans.factory.annotation.Value;
|
|
||||||
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 {
|
|
||||||
private static final Logger LOG = LoggerFactory.getLogger(CustomLabelResolverReactive.class);
|
|
||||||
|
|
||||||
@Value("${label.key-value:}")
|
|
||||||
private String[] keyValues;
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public Map<String, String> resolve(ServerWebExchange exchange) {
|
|
||||||
// rate limit by some request params. such as query params, headers ..
|
|
||||||
return getLabels(keyValues);
|
|
||||||
}
|
|
||||||
|
|
||||||
private Map<String, String> getLabels(String[] keyValues) {
|
|
||||||
Map<String, String> labels = new HashMap<>();
|
|
||||||
for (String kv : keyValues) {
|
|
||||||
String key = kv.substring(0, kv.indexOf(":"));
|
|
||||||
String value = kv.substring(kv.indexOf(":") + 1);
|
|
||||||
labels.put(key, value);
|
|
||||||
}
|
|
||||||
|
|
||||||
LOG.info("Current labels:{}", labels);
|
|
||||||
return labels;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Loading…
Reference in new issue