Browse Source

[Improvement][ApiServer] Traffic limit #5235 (#5307)

* [Improvement][ApiServer] Traffic limit #5235

* code style check
pull/3/MERGE
wenjun 4 years ago committed by GitHub
parent
commit
7843ed4087
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 22
      dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/configuration/AppConfiguration.java
  2. 78
      dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/configuration/TrafficConfiguration.java
  3. 104
      dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/interceptor/RateLimitInterceptor.java
  4. 9
      dolphinscheduler-api/src/main/resources/application-api.properties
  5. 60
      dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/configuration/TrafficConfigurationTest.java
  6. 89
      dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/interceptor/RateLimitInterceptorTest.java
  7. 2
      pom.xml

22
dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/configuration/AppConfiguration.java

@ -19,17 +19,24 @@ package org.apache.dolphinscheduler.api.configuration;
import org.apache.dolphinscheduler.api.interceptor.LocaleChangeInterceptor;
import org.apache.dolphinscheduler.api.interceptor.LoginHandlerInterceptor;
import org.apache.dolphinscheduler.api.interceptor.RateLimitInterceptor;
import java.util.Locale;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.cors.CorsConfiguration;
import org.springframework.web.cors.UrlBasedCorsConfigurationSource;
import org.springframework.web.filter.CorsFilter;
import org.springframework.web.servlet.LocaleResolver;
import org.springframework.web.servlet.config.annotation.*;
import org.springframework.web.servlet.config.annotation.ContentNegotiationConfigurer;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.ResourceHandlerRegistry;
import org.springframework.web.servlet.config.annotation.ViewControllerRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
import org.springframework.web.servlet.i18n.CookieLocaleResolver;
import java.util.Locale;
/**
* application configuration
@ -43,6 +50,9 @@ public class AppConfiguration implements WebMvcConfigurer {
public static final String PATH_PATTERN = "/**";
public static final String LOCALE_LANGUAGE_COOKIE = "language";
@Autowired
private TrafficConfiguration trafficConfiguration;
@Bean
public CorsFilter corsFilter() {
CorsConfiguration config = new CorsConfiguration();
@ -79,10 +89,18 @@ public class AppConfiguration implements WebMvcConfigurer {
return new LocaleChangeInterceptor();
}
@Bean
public RateLimitInterceptor createRateLimitInterceptor() {
return new RateLimitInterceptor(trafficConfiguration);
}
@Override
public void addInterceptors(InterceptorRegistry registry) {
// i18n
registry.addInterceptor(localeChangeInterceptor());
if (trafficConfiguration.isTrafficGlobalControlSwitch() || trafficConfiguration.isTrafficTenantControlSwitch()) {
registry.addInterceptor(createRateLimitInterceptor());
}
registry.addInterceptor(loginInterceptor())
.addPathPatterns(LOGIN_INTERCEPTOR_PATH_PATTERN)
.excludePathPatterns(LOGIN_PATH_PATTERN, REGISTER_PATH_PATTERN,

78
dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/configuration/TrafficConfiguration.java

@ -0,0 +1,78 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* 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 org.apache.dolphinscheduler.api.configuration;
import java.util.Map;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Configuration;
@Configuration
public class TrafficConfiguration {
@Value("${traffic.control.global.switch:false}")
private boolean trafficGlobalControlSwitch;
@Value("${traffic.control.max.global.qps.rate:300}")
private Integer maxGlobalQpsRate;
@Value("${traffic.control.tenant.switch:false}")
private boolean trafficTenantControlSwitch;
@Value("${traffic.control.default.tenant.qps.rate:10}")
private Integer defaultTenantQpsRate;
@Value("#{'${traffic.control.customize.tenant.qps.rate:}'.empty?null:'${traffic.control.customize.tenant.qps.rate:}'}")
private Map<String, Integer> customizeTenantQpsRate;
public boolean isTrafficGlobalControlSwitch() {
return trafficGlobalControlSwitch;
}
public void setTrafficGlobalControlSwitch(boolean trafficGlobalControlSwitch) {
this.trafficGlobalControlSwitch = trafficGlobalControlSwitch;
}
public Integer getMaxGlobalQpsRate() {
return maxGlobalQpsRate;
}
public void setMaxGlobalQpsRate(Integer maxGlobalQpsRate) {
this.maxGlobalQpsRate = maxGlobalQpsRate;
}
public boolean isTrafficTenantControlSwitch() {
return trafficTenantControlSwitch;
}
public void setTrafficTenantControlSwitch(boolean trafficTenantControlSwitch) {
this.trafficTenantControlSwitch = trafficTenantControlSwitch;
}
public Integer getDefaultTenantQpsRate() {
return defaultTenantQpsRate;
}
public void setDefaultTenantQpsRate(Integer defaultTenantQpsRate) {
this.defaultTenantQpsRate = defaultTenantQpsRate;
}
public Map<String, Integer> getCustomizeTenantQpsRate() {
return customizeTenantQpsRate;
}
public void setCustomizeTenantQpsRate(Map<String, Integer> customizeTenantQpsRate) {
this.customizeTenantQpsRate = customizeTenantQpsRate;
}
}

104
dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/interceptor/RateLimitInterceptor.java

@ -0,0 +1,104 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* 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 org.apache.dolphinscheduler.api.interceptor;
import org.apache.dolphinscheduler.api.configuration.TrafficConfiguration;
import org.apache.dolphinscheduler.common.utils.StringUtils;
import org.apache.commons.collections.MapUtils;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.web.servlet.HandlerInterceptor;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
/**
* This interceptor is used to control the traffic, consists with global traffic control and tenant-leve traffic control.
* If the current coming tenant reaches his tenant-level request quota, his request will be reject fast.
* If the current system request number reaches the global request quota, all coming request will be reject fast.
*/
public class RateLimitInterceptor implements HandlerInterceptor {
private static final Logger logger = LoggerFactory.getLogger(RateLimitInterceptor.class);
private TrafficConfiguration trafficConfiguration;
private RateLimiter globalRateLimiter;
private LoadingCache<String, RateLimiter> tenantRateLimiterCache = CacheBuilder.newBuilder()
.maximumSize(100)
.expireAfterAccess(10, TimeUnit.MINUTES)
.build(new CacheLoader<String, RateLimiter>() {
@Override
public RateLimiter load(String token) {
// use tenant customize rate limit
Map<String, Integer> customizeTenantQpsRate = trafficConfiguration.getCustomizeTenantQpsRate();
int tenantQuota = trafficConfiguration.getDefaultTenantQpsRate();
if (MapUtils.isNotEmpty(customizeTenantQpsRate)) {
tenantQuota = customizeTenantQpsRate.getOrDefault(token, trafficConfiguration.getDefaultTenantQpsRate());
}
// use tenant default rate limit
return RateLimiter.create(tenantQuota, 1, TimeUnit.SECONDS);
}
});
@Override
public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws ExecutionException {
// tenant-level rate limit
if (trafficConfiguration.isTrafficTenantControlSwitch()) {
String token = request.getHeader("token");
if (StringUtils.isNotEmpty(token)) {
RateLimiter tenantRateLimiter = tenantRateLimiterCache.get(token);
if (!tenantRateLimiter.tryAcquire()) {
response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
logger.warn("Too many request, reach tenant rate limit, current tenant:{} qps is {}", token, tenantRateLimiter.getRate());
return false;
}
}
}
// global rate limit
if (trafficConfiguration.isTrafficGlobalControlSwitch()) {
if (!globalRateLimiter.tryAcquire()) {
response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
logger.warn("Too many request, reach global rate limit, current qps is {}", globalRateLimiter.getRate());
return false;
}
}
return true;
}
public RateLimitInterceptor(TrafficConfiguration trafficConfiguration) {
this.trafficConfiguration = trafficConfiguration;
if (trafficConfiguration.isTrafficGlobalControlSwitch()) {
this.globalRateLimiter = RateLimiter.create(trafficConfiguration.getMaxGlobalQpsRate(), 1, TimeUnit.SECONDS);
}
}
}

9
dolphinscheduler-api/src/main/resources/application-api.properties

@ -47,6 +47,15 @@ spring.messages.basename=i18n/messages
# Authentication types (supported types: PASSWORD)
security.authentication.type=PASSWORD
# Traffic control, if you turn on this config, the maximum number of request/s will be limited.
# global max request number per second
# default tenant-level max request number
#traffic.control.global.switch=true
#traffic.control.max.global.qps.rate=500
#traffic.control.tenant.switch=true
#traffic.control.default.tenant.qps.rate=10
#traffic.control.customize.tenant.qps.rate={'tenant1':11,'tenant2':20}
#============================================================================
# LDAP Config
# mock ldap server from https://www.forumsys.com/tutorials/integration-how-to/ldap/online-ldap-test-server/

60
dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/configuration/TrafficConfigurationTest.java

@ -0,0 +1,60 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* 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 org.apache.dolphinscheduler.api.configuration;
import org.apache.commons.collections.MapUtils;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;
@RunWith(SpringRunner.class)
@SpringBootTest
public class TrafficConfigurationTest {
@Autowired
private TrafficConfiguration trafficConfiguration;
@Test
public void isTrafficGlobalControlSwitch() {
Assert.assertFalse(trafficConfiguration.isTrafficGlobalControlSwitch());
}
@Test
public void getMaxGlobalQpsLimit() {
Assert.assertEquals(300, (int) trafficConfiguration.getMaxGlobalQpsRate());
}
@Test
public void isTrafficTenantControlSwitch() {
Assert.assertFalse(trafficConfiguration.isTrafficTenantControlSwitch());
}
@Test
public void getDefaultTenantQpsLimit() {
Assert.assertEquals(10, (int) trafficConfiguration.getDefaultTenantQpsRate());
}
@Test
public void getCustomizeTenantQpsRate() {
Assert.assertTrue(MapUtils.isEmpty(trafficConfiguration.getCustomizeTenantQpsRate()));
}
}

89
dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/interceptor/RateLimitInterceptorTest.java

@ -0,0 +1,89 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* 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 org.apache.dolphinscheduler.api.interceptor;
import org.apache.dolphinscheduler.api.configuration.TrafficConfiguration;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.powermock.api.mockito.PowerMockito;
import org.powermock.modules.junit4.PowerMockRunner;
@RunWith(PowerMockRunner.class)
public class RateLimitInterceptorTest {
@Test
public void testPreHandleWithoutControl() throws ExecutionException {
HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
RateLimitInterceptor rateLimitInterceptor = new RateLimitInterceptor(new TrafficConfiguration());
Assert.assertTrue(rateLimitInterceptor.preHandle(request, response, null));
Assert.assertTrue(rateLimitInterceptor.preHandle(request, response, null));
}
@Test
public void testPreHandleWithTenantLevenControl() throws ExecutionException {
TrafficConfiguration trafficConfiguration = new TrafficConfiguration();
trafficConfiguration.setTrafficTenantControlSwitch(true);
Map<String, Integer> map = new HashMap<>();
map.put("tenant1", 2);
map.put("tenant2", 2);
trafficConfiguration.setCustomizeTenantQpsRate(map);
trafficConfiguration.setDefaultTenantQpsRate(4);
RateLimitInterceptor rateLimitInterceptor = new RateLimitInterceptor(trafficConfiguration);
HttpServletRequest tenant1Request = Mockito.mock(HttpServletRequest.class);
HttpServletRequest tenant2Request = Mockito.mock(HttpServletRequest.class);
PowerMockito.when(tenant1Request.getHeader(Mockito.any())).thenReturn("tenant1");
PowerMockito.when(tenant2Request.getHeader(Mockito.any())).thenReturn("tenant2");
HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
for (int i = 0; i < 2; i++) {
rateLimitInterceptor.preHandle(tenant1Request, response, null);
}
Assert.assertFalse(rateLimitInterceptor.preHandle(tenant1Request, response, null));
Assert.assertTrue(rateLimitInterceptor.preHandle(tenant2Request, response, null));
}
@Test
public void testPreHandleWithGlobalControl() throws ExecutionException {
TrafficConfiguration trafficConfiguration = new TrafficConfiguration();
trafficConfiguration.setTrafficTenantControlSwitch(true);
trafficConfiguration.setTrafficGlobalControlSwitch(true);
trafficConfiguration.setMaxGlobalQpsRate(3);
RateLimitInterceptor rateLimitInterceptor = new RateLimitInterceptor(trafficConfiguration);
HttpServletRequest request = Mockito.mock(HttpServletRequest.class);
HttpServletResponse response = Mockito.mock(HttpServletResponse.class);
for (int i = 0; i < 2; i++) {
rateLimitInterceptor.preHandle(request, response, null);
}
Assert.assertFalse(rateLimitInterceptor.preHandle(request, response, null));
}
}

2
pom.xml

@ -795,6 +795,7 @@
<version>${maven-surefire-plugin.version}</version>
<configuration>
<includes>
<include>**/api/configuration/TrafficConfigurationTest.java</include>
<include>**/api/controller/ProcessDefinitionControllerTest.java</include>
<include>**/api/controller/TenantControllerTest.java</include>
<include>**/api/dto/resources/filter/ResourceFilterTest.java</include>
@ -805,6 +806,7 @@
<include>**/api/exceptions/ServiceExceptionTest.java</include>
<include>**/api/interceptor/LocaleChangeInterceptorTest.java</include>
<include>**/api/interceptor/LoginHandlerInterceptorTest.java</include>
<include>**/api/interceptor/RateLimitInterceptorTest.java</include>
<include>**/api/security/impl/pwd/PasswordAuthenticatorTest.java</include>
<include>**/api/security/impl/ldap/LdapAuthenticatorTest.java</include>
<include>**/api/security/SecurityConfigLDAPTest.java</include>

Loading…
Cancel
Save