feature:add PolarisRateLimiterLimitedFallback spi.

pull/1080/head
lingxiao.wu 2 years ago committed by Haotian Zhang
parent b98ee5845b
commit 52689515f5

@ -26,6 +26,7 @@ import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckReactiveFilter;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilter; import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilter;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver; import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver; import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.polaris.client.api.SDKContext; import com.tencent.polaris.client.api.SDKContext;
import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.factory.LimitAPIFactory; import com.tencent.polaris.ratelimit.factory.LimitAPIFactory;
@ -78,9 +79,10 @@ public class PolarisRateLimitAutoConfiguration {
public QuotaCheckServletFilter quotaCheckFilter(LimitAPI limitAPI, public QuotaCheckServletFilter quotaCheckFilter(LimitAPI limitAPI,
@Nullable PolarisRateLimiterLabelServletResolver labelResolver, @Nullable PolarisRateLimiterLabelServletResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties, PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver) { RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
return new QuotaCheckServletFilter(limitAPI, labelResolver, return new QuotaCheckServletFilter(limitAPI, labelResolver,
polarisRateLimitProperties, rateLimitRuleLabelResolver); polarisRateLimitProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
} }
@Bean @Bean
@ -106,9 +108,10 @@ public class PolarisRateLimitAutoConfiguration {
public QuotaCheckReactiveFilter quotaCheckReactiveFilter(LimitAPI limitAPI, public QuotaCheckReactiveFilter quotaCheckReactiveFilter(LimitAPI limitAPI,
@Nullable PolarisRateLimiterLabelReactiveResolver labelResolver, @Nullable PolarisRateLimiterLabelReactiveResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties, PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver) { RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
return new QuotaCheckReactiveFilter(limitAPI, labelResolver, return new QuotaCheckReactiveFilter(limitAPI, labelResolver,
polarisRateLimitProperties, rateLimitRuleLabelResolver); polarisRateLimitProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
} }
} }
} }

@ -22,6 +22,7 @@ import java.nio.charset.StandardCharsets;
import java.time.Duration; import java.time.Duration;
import java.util.HashMap; import java.util.HashMap;
import java.util.Map; import java.util.Map;
import java.util.Objects;
import java.util.Set; import java.util.Set;
import javax.annotation.PostConstruct; import javax.annotation.PostConstruct;
@ -33,6 +34,7 @@ import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties; import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant; import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver; import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils; import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils;
import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils; import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils;
import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.core.LimitAPI;
@ -47,6 +49,7 @@ import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer; import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.MediaType; import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse; import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.lang.Nullable;
import org.springframework.web.server.ServerWebExchange; import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter; import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain; import org.springframework.web.server.WebFilterChain;
@ -70,16 +73,21 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver; private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver;
private final PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
private String rejectTips; private String rejectTips;
public QuotaCheckReactiveFilter(LimitAPI limitAPI, public QuotaCheckReactiveFilter(LimitAPI limitAPI,
PolarisRateLimiterLabelReactiveResolver labelResolver, PolarisRateLimiterLabelReactiveResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties, PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver) { RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
this.limitAPI = limitAPI; this.limitAPI = limitAPI;
this.labelResolver = labelResolver; this.labelResolver = labelResolver;
this.polarisRateLimitProperties = polarisRateLimitProperties; this.polarisRateLimitProperties = polarisRateLimitProperties;
this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver; this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver;
this.polarisRateLimiterLimitedFallback = polarisRateLimiterLimitedFallback;
} }
@PostConstruct @PostConstruct
@ -106,10 +114,19 @@ public class QuotaCheckReactiveFilter implements WebFilter, Ordered {
if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) { if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
ServerHttpResponse response = exchange.getResponse(); ServerHttpResponse response = exchange.getResponse();
response.setRawStatusCode(polarisRateLimitProperties.getRejectHttpCode()); DataBuffer dataBuffer;
response.getHeaders().setContentType(MediaType.TEXT_HTML); if (!Objects.isNull(polarisRateLimiterLimitedFallback)) {
DataBuffer dataBuffer = response.bufferFactory().allocateBuffer() response.setRawStatusCode(polarisRateLimiterLimitedFallback.rejectHttpCode());
.write(rejectTips.getBytes(StandardCharsets.UTF_8)); response.getHeaders().setContentType(polarisRateLimiterLimitedFallback.mediaType());
dataBuffer = response.bufferFactory().allocateBuffer()
.write(polarisRateLimiterLimitedFallback.rejectTips().getBytes(polarisRateLimiterLimitedFallback.charset()));
}
else {
response.setRawStatusCode(polarisRateLimitProperties.getRejectHttpCode());
response.getHeaders().setContentType(MediaType.TEXT_HTML);
dataBuffer = response.bufferFactory().allocateBuffer()
.write(rejectTips.getBytes(StandardCharsets.UTF_8));
}
return response.writeWith(Mono.just(dataBuffer)); return response.writeWith(Mono.just(dataBuffer));
} }
// Unirate // Unirate

@ -22,6 +22,7 @@ import java.io.IOException;
import java.util.Collections; import java.util.Collections;
import java.util.HashMap; import java.util.HashMap;
import java.util.Map; import java.util.Map;
import java.util.Objects;
import java.util.Set; import java.util.Set;
import javax.annotation.PostConstruct; import javax.annotation.PostConstruct;
@ -36,6 +37,7 @@ import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties; import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant; import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver; import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils; import com.tencent.cloud.polaris.ratelimit.utils.QuotaCheckUtils;
import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils; import com.tencent.cloud.polaris.ratelimit.utils.RateLimitUtils;
import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.core.LimitAPI;
@ -46,7 +48,9 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.Order; import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.lang.NonNull; import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.web.filter.OncePerRequestFilter; import org.springframework.web.filter.OncePerRequestFilter;
import static com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant.LABEL_METHOD; import static com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant.LABEL_METHOD;
@ -72,16 +76,20 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver; private final RateLimitRuleLabelResolver rateLimitRuleLabelResolver;
private final PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
private String rejectTips; private String rejectTips;
public QuotaCheckServletFilter(LimitAPI limitAPI, public QuotaCheckServletFilter(LimitAPI limitAPI,
PolarisRateLimiterLabelServletResolver labelResolver, PolarisRateLimiterLabelServletResolver labelResolver,
PolarisRateLimitProperties polarisRateLimitProperties, PolarisRateLimitProperties polarisRateLimitProperties,
RateLimitRuleLabelResolver rateLimitRuleLabelResolver) { RateLimitRuleLabelResolver rateLimitRuleLabelResolver,
@Nullable PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback) {
this.limitAPI = limitAPI; this.limitAPI = limitAPI;
this.labelResolver = labelResolver; this.labelResolver = labelResolver;
this.polarisRateLimitProperties = polarisRateLimitProperties; this.polarisRateLimitProperties = polarisRateLimitProperties;
this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver; this.rateLimitRuleLabelResolver = rateLimitRuleLabelResolver;
this.polarisRateLimiterLimitedFallback = polarisRateLimiterLimitedFallback;
} }
@PostConstruct @PostConstruct
@ -103,9 +111,17 @@ public class QuotaCheckServletFilter extends OncePerRequestFilter {
localNamespace, localService, 1, labels, request.getRequestURI()); localNamespace, localService, 1, labels, request.getRequestURI());
if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) { if (quotaResponse.getCode() == QuotaResultCode.QuotaResultLimited) {
response.setStatus(polarisRateLimitProperties.getRejectHttpCode()); if (!Objects.isNull(polarisRateLimiterLimitedFallback)) {
response.setContentType("text/html;charset=UTF-8"); response.setStatus(polarisRateLimiterLimitedFallback.rejectHttpCode());
response.getWriter().write(rejectTips); String contentType = new MediaType(polarisRateLimiterLimitedFallback.mediaType(), polarisRateLimiterLimitedFallback.charset()).toString();
response.setContentType(contentType);
response.getWriter().write(polarisRateLimiterLimitedFallback.rejectTips());
}
else {
response.setStatus(polarisRateLimitProperties.getRejectHttpCode());
response.setContentType("text/html;charset=UTF-8");
response.getWriter().write(rejectTips);
}
return; return;
} }
// Unirate // Unirate

@ -0,0 +1,66 @@
/*
* 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.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
/**
* PolarisRateLimiterLimitedFallback.
*
* @author Lingxiao.Wu
*/
public interface PolarisRateLimiterLimitedFallback {
/**
* Customized mediaType when polaris rateLimiter limited.
*
* @return Customized mediaType
*/
default MediaType mediaType() {
return MediaType.TEXT_HTML;
}
/**
* Customized charset when polaris rateLimiter limited.
*
* @return Customized charset
*/
default Charset charset() {
return StandardCharsets.UTF_8;
}
/**
* Customized rejectHttpCode when polaris rateLimiter limited.
*
* @return Customized rejectHttpCode
*/
default Integer rejectHttpCode() {
return HttpStatus.TOO_MANY_REQUESTS.value();
}
/**
* Customized rejectTips when polaris rateLimiter limited.
*
* @return Customized rejectTips
*/
String rejectTips();
}

@ -0,0 +1,39 @@
/*
* 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.filter;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import org.springframework.http.MediaType;
/**
* JsonPolarisRateLimiterLimitedFallback.
*
* @author Lingxiao.Wu
*/
public class JsonPolarisRateLimiterLimitedFallback implements PolarisRateLimiterLimitedFallback {
@Override
public MediaType mediaType() {
return MediaType.APPLICATION_JSON;
}
@Override
public String rejectTips() {
return "{errMsg:RejectRequestTips提示消息}";
}
}

@ -31,6 +31,7 @@ import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties; import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant; import com.tencent.cloud.polaris.ratelimit.constant.RateLimitConstant;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver; import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult; import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest; import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest;
@ -79,6 +80,8 @@ public class QuotaCheckReactiveFilterTest {
private final PolarisRateLimiterLabelReactiveResolver labelResolver = private final PolarisRateLimiterLabelReactiveResolver labelResolver =
exchange -> Collections.singletonMap("ReactiveResolver", "ReactiveResolver"); exchange -> Collections.singletonMap("ReactiveResolver", "ReactiveResolver");
private QuotaCheckReactiveFilter quotaCheckReactiveFilter; private QuotaCheckReactiveFilter quotaCheckReactiveFilter;
private QuotaCheckReactiveFilter quotaCheckWithRateLimiterLimitedFallbackReactiveFilter;
private PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
@BeforeClass @BeforeClass
public static void beforeClass() { public static void beforeClass() {
@ -100,6 +103,7 @@ public class QuotaCheckReactiveFilterTest {
@Before @Before
public void setUp() { public void setUp() {
MetadataContext.LOCAL_NAMESPACE = "TEST"; MetadataContext.LOCAL_NAMESPACE = "TEST";
polarisRateLimiterLimitedFallback = new JsonPolarisRateLimiterLimitedFallback();
LimitAPI limitAPI = mock(LimitAPI.class); LimitAPI limitAPI = mock(LimitAPI.class);
when(limitAPI.getQuota(any(QuotaRequest.class))).thenAnswer(invocationOnMock -> { when(limitAPI.getQuota(any(QuotaRequest.class))).thenAnswer(invocationOnMock -> {
@ -123,10 +127,12 @@ public class QuotaCheckReactiveFilterTest {
polarisRateLimitProperties.setRejectHttpCode(419); polarisRateLimitProperties.setRejectHttpCode(419);
RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class); RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.EMPTY_SET); when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.emptySet());
this.quotaCheckReactiveFilter = new QuotaCheckReactiveFilter( this.quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(
limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver); limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver, null);
this.quotaCheckWithRateLimiterLimitedFallbackReactiveFilter = new QuotaCheckReactiveFilter(
limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
} }
@Test @Test
@ -169,12 +175,12 @@ public class QuotaCheckReactiveFilterTest {
throw new RuntimeException("Mock exception."); throw new RuntimeException("Mock exception.");
} }
}; };
quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, exceptionLabelResolver, null, null); quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, exceptionLabelResolver, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange); result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange);
assertThat(result.size()).isEqualTo(0); assertThat(result.size()).isEqualTo(0);
// labelResolver == null // labelResolver == null
quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, null, null, null); quotaCheckReactiveFilter = new QuotaCheckReactiveFilter(null, null, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange); result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckReactiveFilter, exchange);
assertThat(result.size()).isEqualTo(0); assertThat(result.size()).isEqualTo(0);
@ -204,7 +210,9 @@ public class QuotaCheckReactiveFilterTest {
MetadataContext.LOCAL_SERVICE = "TestApp2"; MetadataContext.LOCAL_SERVICE = "TestApp2";
long startTimestamp = System.currentTimeMillis(); long startTimestamp = System.currentTimeMillis();
CountDownLatch countDownLatch = new CountDownLatch(1); CountDownLatch countDownLatch = new CountDownLatch(1);
quotaCheckReactiveFilter.filter(exchange, webFilterChain).subscribe(e -> { }, t -> { }, countDownLatch::countDown); quotaCheckReactiveFilter.filter(exchange, webFilterChain).subscribe(e -> {
}, t -> {
}, countDownLatch::countDown);
try { try {
countDownLatch.await(); countDownLatch.await();
} }
@ -225,8 +233,49 @@ public class QuotaCheckReactiveFilterTest {
quotaCheckReactiveFilter.filter(exchange, webFilterChain); quotaCheckReactiveFilter.filter(exchange, webFilterChain);
} }
@Test
public void polarisRateLimiterLimitedFallbackTest() {
// Create mock WebFilterChain
WebFilterChain webFilterChain = serverWebExchange -> Mono.empty();
// Mock request
MockServerHttpRequest request = MockServerHttpRequest.get("http://localhost:8080/test").build();
ServerWebExchange exchange = MockServerWebExchange.from(request);
quotaCheckWithRateLimiterLimitedFallbackReactiveFilter.init();
// Pass
MetadataContext.LOCAL_SERVICE = "TestApp1";
quotaCheckWithRateLimiterLimitedFallbackReactiveFilter.filter(exchange, webFilterChain);
// Unirate waiting 1000ms
MetadataContext.LOCAL_SERVICE = "TestApp2";
long startTimestamp = System.currentTimeMillis();
CountDownLatch countDownLatch = new CountDownLatch(1);
quotaCheckWithRateLimiterLimitedFallbackReactiveFilter.filter(exchange, webFilterChain).subscribe(e -> {
}, t -> {
}, countDownLatch::countDown);
try {
countDownLatch.await();
}
catch (InterruptedException e) {
fail("Exception encountered.", e);
}
assertThat(System.currentTimeMillis() - startTimestamp).isGreaterThanOrEqualTo(1000L);
// Rate limited
MetadataContext.LOCAL_SERVICE = "TestApp3";
quotaCheckWithRateLimiterLimitedFallbackReactiveFilter.filter(exchange, webFilterChain);
ServerHttpResponse response = exchange.getResponse();
assertThat(response.getRawStatusCode()).isEqualTo(polarisRateLimiterLimitedFallback.rejectHttpCode());
assertThat(response.getHeaders().getContentType()).isEqualTo(polarisRateLimiterLimitedFallback.mediaType());
// Exception
MetadataContext.LOCAL_SERVICE = "TestApp4";
quotaCheckWithRateLimiterLimitedFallbackReactiveFilter.filter(exchange, webFilterChain);
}
@SpringBootApplication @SpringBootApplication
protected static class TestApplication { protected static class TestApplication {
} }
} }

@ -34,6 +34,7 @@ import com.tencent.cloud.common.util.expresstion.SpringWebExpressionLabelUtils;
import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver; import com.tencent.cloud.polaris.ratelimit.RateLimitRuleLabelResolver;
import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties; import com.tencent.cloud.polaris.ratelimit.config.PolarisRateLimitProperties;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver; import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult; import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
import com.tencent.polaris.ratelimit.api.core.LimitAPI; import com.tencent.polaris.ratelimit.api.core.LimitAPI;
import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest; import com.tencent.polaris.ratelimit.api.rpc.QuotaRequest;
@ -49,6 +50,7 @@ import org.mockito.junit.MockitoJUnitRunner;
import org.springframework.boot.autoconfigure.SpringBootApplication; import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.test.context.SpringBootTest; import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.http.MediaType;
import org.springframework.mock.web.MockHttpServletRequest; import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.mock.web.MockHttpServletResponse; import org.springframework.mock.web.MockHttpServletResponse;
import org.springframework.web.server.ServerWebExchange; import org.springframework.web.server.ServerWebExchange;
@ -79,6 +81,8 @@ public class QuotaCheckServletFilterTest {
exchange -> Collections.singletonMap("ServletResolver", "ServletResolver"); exchange -> Collections.singletonMap("ServletResolver", "ServletResolver");
private QuotaCheckServletFilter quotaCheckServletFilter; private QuotaCheckServletFilter quotaCheckServletFilter;
private QuotaCheckServletFilter quotaCheckWithHtmlRejectTipsServletFilter; private QuotaCheckServletFilter quotaCheckWithHtmlRejectTipsServletFilter;
private QuotaCheckServletFilter quotaCheckWithRateLimiterLimitedFallbackFilter;
private PolarisRateLimiterLimitedFallback polarisRateLimiterLimitedFallback;
@BeforeClass @BeforeClass
public static void beforeClass() { public static void beforeClass() {
@ -129,9 +133,12 @@ public class QuotaCheckServletFilterTest {
RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class); RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.emptySet()); when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.emptySet());
this.quotaCheckServletFilter = new QuotaCheckServletFilter(limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver); this.quotaCheckServletFilter = new QuotaCheckServletFilter(limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver, null);
this.quotaCheckWithHtmlRejectTipsServletFilter = new QuotaCheckServletFilter( this.quotaCheckWithHtmlRejectTipsServletFilter = new QuotaCheckServletFilter(
limitAPI, labelResolver, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleLabelResolver); limitAPI, labelResolver, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleLabelResolver, null);
polarisRateLimiterLimitedFallback = new JsonPolarisRateLimiterLimitedFallback();
this.quotaCheckWithRateLimiterLimitedFallbackFilter = new QuotaCheckServletFilter(
limitAPI, labelResolver, polarisRateLimitWithHtmlRejectTipsProperties, rateLimitRuleLabelResolver, polarisRateLimiterLimitedFallback);
} }
@Test @Test
@ -154,6 +161,7 @@ public class QuotaCheckServletFilterTest {
catch (NoSuchFieldException | IllegalAccessException e) { catch (NoSuchFieldException | IllegalAccessException e) {
fail("Exception encountered.", e); fail("Exception encountered.", e);
} }
quotaCheckWithRateLimiterLimitedFallbackFilter.init();
} }
@Test @Test
@ -174,12 +182,12 @@ public class QuotaCheckServletFilterTest {
PolarisRateLimiterLabelServletResolver exceptionLabelResolver = request1 -> { PolarisRateLimiterLabelServletResolver exceptionLabelResolver = request1 -> {
throw new RuntimeException("Mock exception."); throw new RuntimeException("Mock exception.");
}; };
quotaCheckServletFilter = new QuotaCheckServletFilter(null, exceptionLabelResolver, null, null); quotaCheckServletFilter = new QuotaCheckServletFilter(null, exceptionLabelResolver, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request); result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request);
assertThat(result.size()).isEqualTo(0); assertThat(result.size()).isEqualTo(0);
// labelResolver == null // labelResolver == null
quotaCheckServletFilter = new QuotaCheckServletFilter(null, null, null, null); quotaCheckServletFilter = new QuotaCheckServletFilter(null, null, null, null, null);
result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request); result = (Map<String, String>) getCustomResolvedLabels.invoke(quotaCheckServletFilter, request);
assertThat(result.size()).isEqualTo(0); assertThat(result.size()).isEqualTo(0);
@ -232,9 +240,46 @@ public class QuotaCheckServletFilterTest {
fail("Exception encountered.", e); fail("Exception encountered.", e);
} }
} }
@Test
public void polarisRateLimiterLimitedFallbackTest() {
// Create mock FilterChain
FilterChain filterChain = (servletRequest, servletResponse) -> {
};
// Mock request
MockHttpServletRequest request = new MockHttpServletRequest();
MockHttpServletResponse response = new MockHttpServletResponse();
quotaCheckWithRateLimiterLimitedFallbackFilter.init();
try {
// Pass
MetadataContext.LOCAL_SERVICE = "TestApp1";
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
// Unirate waiting 1000ms
MetadataContext.LOCAL_SERVICE = "TestApp2";
long startTimestamp = System.currentTimeMillis();
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
assertThat(System.currentTimeMillis() - startTimestamp).isGreaterThanOrEqualTo(1000L);
// Rate limited
MetadataContext.LOCAL_SERVICE = "TestApp3";
String contentType = new MediaType(polarisRateLimiterLimitedFallback.mediaType(), polarisRateLimiterLimitedFallback.charset()).toString();
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
assertThat(response.getStatus()).isEqualTo(polarisRateLimiterLimitedFallback.rejectHttpCode());
assertThat(response.getContentAsString()).isEqualTo(polarisRateLimiterLimitedFallback.rejectTips());
assertThat(response.getContentType()).isEqualTo(contentType);
// Exception
MetadataContext.LOCAL_SERVICE = "TestApp4";
quotaCheckWithRateLimiterLimitedFallbackFilter.doFilterInternal(request, response, filterChain);
}
catch (ServletException | IOException e) {
fail("Exception encountered.", e);
}
}
@SpringBootApplication @SpringBootApplication
protected static class TestApplication { protected static class TestApplication {
} }
} }

@ -0,0 +1,41 @@
/*
* 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 com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLimitedFallback;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
/**
* JsonPolarisRateLimiterLimitedFallback.
*
* @author Lingxiao.Wu
*/
@Component
public class JsonPolarisRateLimiterLimitedFallback implements PolarisRateLimiterLimitedFallback {
@Override
public MediaType mediaType() {
return MediaType.APPLICATION_JSON;
}
@Override
public String rejectTips() {
return "{\"success\":false,\"code\":429,\"msg\":\"RejectRequestTips提示消息\",\"timestamp\":1673415310479}";
}
}
Loading…
Cancel
Save