feat(openapi): 重构 ConsumerToken 限流功能

apollo-common/src/main/java/com/ctrip/framework/apollo/common/exception/BadRequestException.java

@@ -37,6 +37,10 @@ public static BadRequestException orgIdIsBlank() {
     return new BadRequestException("orgId can not be blank");
   }
 
+  public static BadRequestException rateLimitIsInvalid() {
+    return new BadRequestException("Ratelimit must be greater than 1");
+  }
+
   public static BadRequestException itemAlreadyExists(String itemKey) {
     return new BadRequestException("item already exists for itemKey:%s", itemKey);
   }

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/entity/ConsumerToken.java

@@ -18,6 +18,7 @@
 
 import com.ctrip.framework.apollo.common.entity.BaseEntity;
 
+import javax.validation.constraints.PositiveOrZero;
 import org.hibernate.annotations.SQLDelete;
 import org.hibernate.annotations.Where;
 
@@ -41,8 +42,9 @@ public class ConsumerToken extends BaseEntity {
   @Column(name = "`Token`", nullable = false)
   private String token;
 
-  @Column(name = "LimitCount")
-  private Integer limitCount;
+  @PositiveOrZero
+  @Column(name = "`RateLimit`", nullable = false)
+  private Integer rateLimit;
 
   @Column(name = "`Expires`", nullable = false)
   private Date expires;
@@ -63,12 +65,12 @@ public void setToken(String token) {
     this.token = token;
   }
 
-  public Integer getLimitCount() {
-    return limitCount;
+  public Integer getRateLimit() {
+    return rateLimit;
   }
 
-  public void setLimitCount(Integer limitCount) {
-    this.limitCount = limitCount;
+  public void setRateLimit(Integer rateLimit) {
+    this.rateLimit = rateLimit;
   }
 
   public Date getExpires() {
@@ -82,7 +84,7 @@ public void setExpires(Date expires) {
   @Override
   public String toString() {
     return toStringHelper().add("consumerId", consumerId).add("token", token)
-        .add("limitCount", limitCount)
+        .add("rateLimit", rateLimit)
         .add("expires", expires).toString();
   }
 }

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/filter/ConsumerAuthenticationFilter.java

@@ -24,6 +24,7 @@
 import com.google.common.cache.CacheBuilder;
 import com.google.common.util.concurrent.RateLimiter;
 import java.io.IOException;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import javax.servlet.Filter;
 import javax.servlet.FilterChain;
@@ -52,6 +53,8 @@ public class ConsumerAuthenticationFilter implements Filter {
   private static final int WARMUP_MILLIS = 1000; // ms
   private static final int RATE_LIMITER_CACHE_MAX_SIZE = 50000;
 
+  private static final int TOO_MANY_REQUESTS = 429;
+
   private static final Cache<String, ImmutablePair<Long, RateLimiter>> LIMITER = CacheBuilder.newBuilder()
       .expireAfterWrite(1, TimeUnit.DAYS)
       .maximumSize(RATE_LIMITER_CACHE_MAX_SIZE).build();
@@ -76,21 +79,18 @@ public void doFilter(ServletRequest req, ServletResponse resp, FilterChain chain
     String token = request.getHeader(HttpHeaders.AUTHORIZATION);
     ConsumerToken consumerToken = consumerAuthUtil.getConsumerToken(token);
 
-    if (null == consumerToken || consumerToken.getConsumerId() <= 0) {
+    if (null == consumerToken) {
       response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Unauthorized");
       return;
     }
 
-    Integer limitCount = consumerToken.getLimitCount();
-    if (limitCount == null) {
-      limitCount = 0;
-    }
-    if (portalConfig.isOpenApiLimitEnabled() && limitCount > 0) {
+    Integer rateLimit = consumerToken.getRateLimit();
+    if (null != rateLimit && rateLimit > 0) {
       try {
-        ImmutablePair<Long, RateLimiter> rateLimiterPair = getOrCreateRateLimiterPair(token, limitCount);
+        ImmutablePair<Long, RateLimiter> rateLimiterPair = getOrCreateRateLimiterPair(token, rateLimit);
         long warmupToMillis = rateLimiterPair.getLeft() + WARMUP_MILLIS;
         if (System.currentTimeMillis() > warmupToMillis && !rateLimiterPair.getRight().tryAcquire()) {
-          response.sendError(HttpServletResponse.SC_FORBIDDEN, "Too many call requests, the flow is limited");
+          response.sendError(TOO_MANY_REQUESTS, "Too Many Requests, the flow is limited");
           return;
         }
       } catch (Exception e) {
@@ -113,12 +113,12 @@ public void destroy() {
   }
 
   private ImmutablePair<Long, RateLimiter> getOrCreateRateLimiterPair(String key, Integer limitCount) {
-    ImmutablePair<Long, RateLimiter> rateLimiterPair = LIMITER.getIfPresent(key);
-    if (rateLimiterPair == null) {
-      rateLimiterPair = ImmutablePair.of(System.currentTimeMillis(), RateLimiter.create(limitCount));
-      LIMITER.put(key, rateLimiterPair);
+    try{
+      return LIMITER.get(key, () ->
+          ImmutablePair.of(System.currentTimeMillis(), RateLimiter.create(limitCount)));
+    } catch (ExecutionException e) {
+      throw new RuntimeException("Failed to create rate limiter", e);
     }
-    return rateLimiterPair;
   }
 
 }

apollo-portal/src/main/java/com/ctrip/framework/apollo/openapi/service/ConsumerService.java

@@ -120,10 +120,10 @@ public Consumer createConsumer(Consumer consumer) {
     return consumerRepository.save(consumer);
   }
 
-  public ConsumerToken generateAndSaveConsumerToken(Consumer consumer, Date expires) {
+  public ConsumerToken generateAndSaveConsumerToken(Consumer consumer, Integer rateLimit, Date expires) {
     Preconditions.checkArgument(consumer != null, "Consumer can not be null");
 
-    ConsumerToken consumerToken = generateConsumerToken(consumer, expires);
+    ConsumerToken consumerToken = generateConsumerToken(consumer, rateLimit, expires);
     consumerToken.setId(0);
 
     return consumerTokenRepository.save(consumerToken);
@@ -314,20 +314,17 @@ public void createConsumerAudits(Iterable<ConsumerAudit> consumerAudits) {
   @Transactional
   public ConsumerToken createConsumerToken(ConsumerToken entity) {
     entity.setId(0); //for protection
-    if (entity.getLimitCount() <= 0) {
-      entity.setLimitCount(portalConfig.openApiLimitCount());
-    }
     return consumerTokenRepository.save(entity);
   }
 
-  private ConsumerToken generateConsumerToken(Consumer consumer, Date expires) {
+  private ConsumerToken generateConsumerToken(Consumer consumer, Integer rateLimit, Date expires) {
     long consumerId = consumer.getId();
     String createdBy = userInfoHolder.getUser().getUserId();
     Date createdTime = new Date();
 
     ConsumerToken consumerToken = new ConsumerToken();
     consumerToken.setConsumerId(consumerId);
-    consumerToken.setLimitCount(portalConfig.openApiLimitCount());
+    consumerToken.setRateLimit(rateLimit);
     consumerToken.setExpires(expires);
     consumerToken.setDataChangeCreatedBy(createdBy);
     consumerToken.setDataChangeCreatedTime(createdTime);

apollo-portal/src/main/java/com/ctrip/framework/apollo/portal/component/config/PortalConfig.java

@@ -242,14 +242,6 @@ public String consumerTokenSalt() {
     return getValue("consumer.token.salt", "apollo-portal");
   }
 
-  public int openApiLimitCount() {
-    return getIntProperty("open.api.limit.count", 20);
-  }
-
-  public boolean isOpenApiLimitEnabled() {
-    return getBooleanProperty("open.api.limit.enabled", false);
-  }
-
   public boolean isEmailEnabled() {
     return getBooleanProperty("email.enabled", false);
   }

apollo-portal/src/main/java/com/ctrip/framework/apollo/portal/controller/ConsumerController.java

@@ -81,13 +81,21 @@ public ConsumerInfo create(
       throw BadRequestException.orgIdIsBlank();
     }
 
+    if (requestVO.isRateLimitEenabled()) {
+      if (requestVO.getRateLimit() <= 0) {
+        throw BadRequestException.rateLimitIsInvalid();
+      }
+    } else {
+      requestVO.setRateLimit(0);
+    }
+
     Consumer createdConsumer = consumerService.createConsumer(convertToConsumer(requestVO));
 
     if (Objects.isNull(expires)) {
       expires = DEFAULT_EXPIRES;
     }
 
-    ConsumerToken consumerToken = consumerService.generateAndSaveConsumerToken(createdConsumer, expires);
+    ConsumerToken consumerToken = consumerService.generateAndSaveConsumerToken(createdConsumer, requestVO.getRateLimit(), expires);
     if (requestVO.isAllowCreateApplication()) {
       consumerService.assignCreateApplicationRoleToConsumer(consumerToken.getToken());
     }
@@ -127,7 +135,7 @@ public List<ConsumerRole> assignNamespaceRoleToConsumer(
     if (StringUtils.isEmpty(namespaceName)) {
       throw new BadRequestException("Params(NamespaceName) can not be empty.");
     }
-    if (null != envs){
+    if (null != envs) {
       String[] envArray = envs.split(",");
       List<String> envList = Lists.newArrayList();
       // validate env parameter
@@ -156,7 +164,7 @@ public List<ConsumerRole> assignNamespaceRoleToConsumer(
 
   @GetMapping("/consumers")
   @PreAuthorize(value = "@permissionValidator.isSuperAdmin()")
-  public List<ConsumerInfo> getConsumerList(Pageable page){
+  public List<ConsumerInfo> getConsumerList(Pageable page) {
     return consumerService.findConsumerInfoList(page);
   }
 

apollo-portal/src/main/java/com/ctrip/framework/apollo/portal/entity/vo/consumer/ConsumerCreateRequestVO.java

@@ -26,6 +26,8 @@ public class ConsumerCreateRequestVO {
   private String orgId;
   private String orgName;
   private String ownerName;
+  private boolean rateLimitEenabled;
+  private int rateLimit;
 
   public String getAppId() {
     return appId;
@@ -75,4 +77,20 @@ public void setOwnerName(String ownerName) {
     this.ownerName = ownerName;
   }
 
+  public boolean isRateLimitEenabled() {
+    return rateLimitEenabled;
+  }
+
+  public void setRateLimitEenabled(boolean rateLimitEenabled) {
+    this.rateLimitEenabled = rateLimitEenabled;
+  }
+
+  public int getRateLimit() {
+    return rateLimit;
+  }
+
+  public void setRateLimit(int rateLimit) {
+    this.rateLimit = rateLimit;
+  }
+
 }

apollo-portal/src/main/resources/static/i18n/en.json

@@ -654,6 +654,11 @@
   "Open.Manage.Consumer.AllowCreateApplicationTips": "（Allow third-party applications to create apps and grant them app administrator privileges.",
   "Open.Manage.Consumer.AllowCreateApplication.No": "no",
   "Open.Manage.Consumer.AllowCreateApplication.Yes": "yes",
+  "Open.Manage.Consumer.RateLimit.Enabled": "Whether to enable current limit",
+  "Open.Manage.Consumer.RateLimit.Enabled.Tips": "(After enabling this feature, when third-party applications publish configurations on Apollo, their traffic will be controlled according to the configured QPS limit)",
+  "Open.Manage.Consumer.RateLimitValue": "Current limiting QPS",
+  "Open.Manage.Consumer.RateLimitValueTips": "(Unit: times/second, for example: 100 means that the configuration is published at most 100 times per second)",
+  "Open.Manage.Consumer.RateLimitValue.Error": "The minimum current limiting QPS is 1",
   "Namespace.Role.Title": "Permission Management",
   "Namespace.Role.GrantModifyTo": "Permission to edit",
   "Namespace.Role.GrantModifyTo2": "(Can edit the configuration)",

apollo-portal/src/main/resources/static/i18n/zh-CN.json

@@ -654,6 +654,11 @@
   "Open.Manage.Consumer.AllowCreateApplicationTips": "（允许第三方应用创建app，并且对创建出的app，拥有应用管理员的权限）",
   "Open.Manage.Consumer.AllowCreateApplication.No": "否",
   "Open.Manage.Consumer.AllowCreateApplication.Yes": "是",
+  "Open.Manage.Consumer.RateLimit.Enabled": "是否启用限流",
+  "Open.Manage.Consumer.RateLimit.Enabled.Tips": "（开启后，第三方应用在 Apollo 上发布配置时，会根据配置的 QPS 限制，控制其流量）",
+  "Open.Manage.Consumer.RateLimitValue": "限流QPS",
+  "Open.Manage.Consumer.RateLimitValueTips": "（单位:次/秒，例如: 100 表示每秒最多发布 100 次配置）",
+  "Open.Manage.Consumer.RateLimitValue.Error": "限流QPS最小为1",
   "Namespace.Role.Title": "权限管理",
   "Namespace.Role.GrantModifyTo": "修改权",
   "Namespace.Role.GrantModifyTo2": "（可以修改配置）",

apollo-portal/src/main/resources/static/open/add-consumer.html

@@ -78,6 +78,33 @@ <h5>{{'Open.Manage.CreateThirdApp' | translate }}
                             </div>
                         </div>
 
+                        <div class="form-group">
+                            <label class="col-sm-2 control-label">
+                                {{ 'Open.Manage.Consumer.RateLimit.Enabled' | translate }}
+                            </label>
+                            <div class="col-sm-3">
+                                <input type="checkbox"
+                                       ng-model="consumer.rateLimitEenabled"
+                                       name="rateLimitEenabled"
+                                       ng-change="toggleRateLimitEenabledInput()"
+                                />
+                                <small>{{ 'Open.Manage.Consumer.RateLimit.Enabled.Tips' | translate }}</small>
+                            </div>
+                        </div>
+
+                        <div class="form-group" ng-show="consumer.rateLimitEenabled">
+                            <label class="col-sm-2 control-label">
+                                {{ 'Open.Manage.Consumer.RateLimitValue' | translate }}
+                            </label>
+                            <div class="col-sm-3">
+                                <input type="number"
+                                       ng-model="consumer.rateLimit"
+                                       name="rateLimit"
+                                />
+                                <small>{{'Open.Manage.Consumer.RateLimitValueTips' | translate }}</small>
+                            </div>
+                        </div>
+
                         <div class="form-group">
                             <label class="col-sm-2 control-label">
                                 <apollorequiredfield></apollorequiredfield>
@@ -252,4 +279,4 @@ <h4>{{'Common.IsRootUser' | translate }}</h4>
     <script type="application/javascript" src="../scripts/controller/open/OpenManageController.js"></script>
 </body>
 
-</html>
+</html>

apollo-portal/src/main/resources/static/scripts/controller/open/OpenManageController.js

@@ -41,6 +41,11 @@ function OpenManageController($scope, $translate, toastr, AppUtil, OrganizationS
     $scope.preDeleteConsumer = preDeleteConsumer;
     $scope.deleteConsumer = deleteConsumer;
     $scope.preGrantPermission = preGrantPermission;
+    $scope.toggleRateLimitEenabledInput = function() {
+        if (!$scope.consumer.rateLimitEenabled) {
+            $scope.consumer.rateLimit = 0;
+        }
+    };
 
     function init() {
         initOrganization();
@@ -163,6 +168,17 @@ function OpenManageController($scope, $translate, toastr, AppUtil, OrganizationS
             $scope.submitBtnDisabled = false;
             return;
         }
+
+        if ($scope.consumer.rateLimitEenabled) {
+            if ($scope.consumer.rateLimit < 1) {
+                toastr.warning($translate.instant('Open.Manage.Consumer.RateLimitValue.Error'));
+                $scope.submitBtnDisabled = false;
+                return;
+            }
+        } else {
+            $scope.consumer.rateLimit = 0;
+        }
+
         var selectedOrg = $orgWidget.select2('data')[0];
 
         if (!selectedOrg.id) {

apollo-portal/src/test/java/com/ctrip/framework/apollo/SkipAuthorizationConfiguration.java

@@ -54,7 +54,7 @@ public ConsumerAuthUtil consumerAuthUtil() {
 
     ConsumerToken someConsumerToken = new ConsumerToken();
     someConsumerToken.setConsumerId(1L);
-    someConsumerToken.setLimitCount(20);
+    someConsumerToken.setRateLimit(20);
     when(mock.getConsumerToken(any())).thenReturn(someConsumerToken);
     return mock;
   }

apollo-portal/src/test/java/com/ctrip/framework/apollo/openapi/filter/ConsumerAuthenticationFilterTest.java

@@ -53,6 +53,9 @@
  */
 @RunWith(MockitoJUnitRunner.class)
 public class ConsumerAuthenticationFilterTest {
+
+  private static final int TOO_MANY_REQUESTS = 429;
+
   private ConsumerAuthenticationFilter authenticationFilter;
   @Mock
   private ConsumerAuthUtil consumerAuthUtil;
@@ -161,7 +164,7 @@ public void testRateLimitPartFailure() throws Exception {
     int leastTimes = qps * durationInSeconds;
     int mostTimes = (qps + 1) * durationInSeconds;
 
-    verify(response, atLeastOnce()).sendError(eq(HttpServletResponse.SC_FORBIDDEN), anyString());
+    verify(response, atLeastOnce()).sendError(eq(TOO_MANY_REQUESTS), anyString());
 
     verify(consumerAuthUtil, atLeast(leastTimes)).storeConsumerId(request, someConsumerId);
     verify(consumerAuthUtil, atMost(mostTimes)).storeConsumerId(request, someConsumerId);
@@ -176,11 +179,10 @@ public void testRateLimitPartFailure() throws Exception {
   private void setupRateLimitMocks(String someToken, Long someConsumerId, int qps) {
     ConsumerToken someConsumerToken = new ConsumerToken();
     someConsumerToken.setConsumerId(someConsumerId);
-    someConsumerToken.setLimitCount(qps);
+    someConsumerToken.setRateLimit(qps);
 
     when(request.getHeader(HttpHeaders.AUTHORIZATION)).thenReturn(someToken);
     when(consumerAuthUtil.getConsumerToken(someToken)).thenReturn(someConsumerToken);
-    when(portalConfig.isOpenApiLimitEnabled()).thenReturn(true);
   }