添加Redis分布式限流方式实现

Author: hengboy

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/java/org/minbox/framework/api/boot/plugin/rate/limiter/ApiBootRateLimiter.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright [2019] [恒宇少年 - 于起宇]
+ *
+ *      Licensed 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.minbox.framework.api.boot.plugin.rate.limiter;
+
+/**
+ * ApiBoot RateLimiter
+ *
+ * @author：恒宇少年 - 于起宇
+ * <p>
+ * DateTime：2019-05-05 17:12
+ * Blog：http://blog.yuqiyu.com
+ * WebSite：http://www.jianshu.com/u/092df3f77bca
+ * Gitee：https://gitee.com/hengboy
+ * GitHub：https://github.com/hengboy
+ */
+public interface ApiBootRateLimiter {
+    /**
+     * Attempt to obtain a request current limit token
+     *
+     * @param QPS        queries per second
+     * @param requestUri request uri
+     * @return true : allow access to
+     */
+    boolean tryAcquire(Double QPS, String requestUri);
+
+    class Response {
+        private final boolean allowed;
+        private final long tokensRemaining;
+
+        public Response(boolean allowed, long tokensRemaining) {
+            this.allowed = allowed;
+            this.tokensRemaining = tokensRemaining;
+        }
+
+        public boolean isAllowed() {
+            return allowed;
+        }
+
+        @Override
+        public String toString() {
+            return "Response{" +
+                    "allowed=" + allowed +
+                    ", tokensRemaining=" + tokensRemaining +
+                    '}';
+        }
+    }
+}

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/java/org/minbox/framework/api/boot/plugin/rate/limiter/ApiBootRateLimiterConfiguration.java

@@ -54,7 +54,7 @@ public ApiBootRateLimiterConfiguration(RateLimiterConfig rateLimiterConfig, Hand
         this.rateLimiterConfig = rateLimiterConfig;
         this.handlerInterceptor = handlerInterceptor;
         Assert.notNull(rateLimiterConfig, "RateLimiterConfig can not be null.");
-        Assert.notNull(rateLimiterConfig, "HandlerInterceptor can not be null.");
+        Assert.notNull(handlerInterceptor, "HandlerInterceptor can not be null.");
     }
 
     @Override

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/java/org/minbox/framework/api/boot/plugin/rate/limiter/handler/ApiBootDefaultRateLimiterInterceptorHandler.java

@@ -17,10 +17,12 @@
 
 package org.minbox.framework.api.boot.plugin.rate.limiter.handler;
 
+import org.minbox.framework.api.boot.plugin.rate.limiter.ApiBootRateLimiter;
 import org.minbox.framework.api.boot.plugin.rate.limiter.annotation.RateLimiter;
-import org.minbox.framework.api.boot.plugin.rate.limiter.context.ApiBootRateLimiterContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
 import org.springframework.web.method.HandlerMethod;
 import org.springframework.web.servlet.HandlerInterceptor;
 
@@ -43,6 +45,16 @@ public class ApiBootDefaultRateLimiterInterceptorHandler implements HandlerInter
      * logger instance
      */
     static Logger logger = LoggerFactory.getLogger(ApiBootDefaultRateLimiterInterceptorHandler.class);
+    /**
+     * ApiBoot RateLimiter
+     */
+    private ApiBootRateLimiter apiBootRateLimiter;
+
+    public ApiBootDefaultRateLimiterInterceptorHandler(ApiBootRateLimiter apiBootRateLimiter) {
+        this.apiBootRateLimiter = apiBootRateLimiter;
+        Assert.notNull(apiBootRateLimiter, "No ApiBootRateLimiter implementation class instance.");
+        logger.info("ApiBootDefaultRateLimiterInterceptorHandler load complete.");
+    }
 
     /**
      * Executing traffic restrictions
@@ -59,12 +71,14 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
             HandlerMethod handlerMethod = (HandlerMethod) handler;
             // get method declared RateLimiter
             RateLimiter rateLimiterAnnotation = handlerMethod.getMethodAnnotation(RateLimiter.class);
-            com.google.common.util.concurrent.RateLimiter rateLimiter = ApiBootRateLimiterContext.cacheRateLimiter(request.getRequestURI(), rateLimiterAnnotation.QPS());
-            double acquire = rateLimiter.acquire();
-            logger.debug("ApiBoot rate limiter acquire：{}", acquire);
+            // if dont't set RateLimiter
+            if (ObjectUtils.isEmpty(rateLimiterAnnotation)) {
+                return true;
+            }
+            return apiBootRateLimiter.tryAcquire(rateLimiterAnnotation.QPS(), request.getRequestURI());
         } catch (Exception e) {
+            logger.error("Current Limiting Request Encountered Exception.", e);
             return false;
         }
-        return true;
     }
 }

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/java/org/minbox/framework/api/boot/plugin/rate/limiter/support/AbstractRateLimiter.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright [2019] [恒宇少年 - 于起宇]
+ *
+ *      Licensed 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.minbox.framework.api.boot.plugin.rate.limiter.support;
+
+import org.minbox.framework.api.boot.plugin.rate.limiter.ApiBootRateLimiter;
+
+/**
+ * ApiBoot RateLimiter Abstract Support
+ *
+ * @author：恒宇少年 - 于起宇
+ * <p>
+ * DateTime：2019-05-05 17:19
+ * Blog：http://blog.yuqiyu.com
+ * WebSite：http://www.jianshu.com/u/092df3f77bca
+ * Gitee：https://gitee.com/hengboy
+ * GitHub：https://github.com/hengboy
+ */
+public abstract class AbstractRateLimiter implements ApiBootRateLimiter {
+
+}

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/java/org/minbox/framework/api/boot/plugin/rate/limiter/support/GoogleGuavaRateLimiter.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright [2019] [恒宇少年 - 于起宇]
+ *
+ *      Licensed 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.minbox.framework.api.boot.plugin.rate.limiter.support;
+
+import org.minbox.framework.api.boot.plugin.rate.limiter.context.ApiBootRateLimiterContext;
+
+/**
+ * google guava rate limiter support
+ *
+ * @author：恒宇少年 - 于起宇
+ * <p>
+ * DateTime：2019-05-05 17:20
+ * Blog：http://blog.yuqiyu.com
+ * WebSite：http://www.jianshu.com/u/092df3f77bca
+ * Gitee：https://gitee.com/hengboy
+ * GitHub：https://github.com/hengboy
+ */
+public class GoogleGuavaRateLimiter extends AbstractRateLimiter {
+    /**
+     * google guava away
+     *
+     * @param QPS        queries per second
+     * @param requestUri request uri
+     * @return true : allow access to
+     */
+    @Override
+    public boolean tryAcquire(Double QPS, String requestUri) {
+        com.google.common.util.concurrent.RateLimiter rateLimiter = ApiBootRateLimiterContext.cacheRateLimiter(requestUri, QPS);
+        return rateLimiter.tryAcquire();
+    }
+}

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/java/org/minbox/framework/api/boot/plugin/rate/limiter/support/RedisLuaRateLimiter.java

@@ -0,0 +1,141 @@
+/*
+ * Copyright [2019] [恒宇少年 - 于起宇]
+ *
+ *      Licensed 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.minbox.framework.api.boot.plugin.rate.limiter.support;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.core.io.ClassPathResource;
+import org.springframework.data.redis.core.RedisTemplate;
+import org.springframework.data.redis.core.script.DefaultRedisScript;
+import org.springframework.data.redis.core.script.RedisScript;
+import org.springframework.data.redis.serializer.RedisSerializer;
+import org.springframework.data.redis.serializer.StringRedisSerializer;
+import org.springframework.scripting.support.ResourceScriptSource;
+import org.springframework.util.Assert;
+
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * redis lua rate limiter support
+ *
+ * @author：恒宇少年 - 于起宇
+ * <p>
+ * DateTime：2019-05-05 17:21
+ * Blog：http://blog.yuqiyu.com
+ * WebSite：http://www.jianshu.com/u/092df3f77bca
+ * Gitee：https://gitee.com/hengboy
+ * GitHub：https://github.com/hengboy
+ */
+public class RedisLuaRateLimiter extends AbstractRateLimiter {
+    /**
+     * logger instance
+     */
+    static Logger logger = LoggerFactory.getLogger(RedisLuaRateLimiter.class);
+
+    /**
+     * Redis Script file name.
+     */
+    private static final String QPS_LUA_PATH = "META-INF/scripts/qps-rate-limiter.lua";
+    /**
+     * redis template
+     */
+    private RedisTemplate redisTemplate;
+    /**
+     * Redis Script.
+     */
+    private RedisScript<List<Long>> redisScript;
+
+    public RedisLuaRateLimiter(RedisTemplate redisTemplate) {
+        this.redisTemplate = redisTemplate;
+        this.redisScript = getRedisScript();
+
+        Assert.notNull(redisTemplate, "No RedisTemplate implementation class was found.");
+        Assert.notNull(redisScript, "Unable to load Lua script.");
+    }
+
+    /**
+     * redis lua away
+     * Processing traffic restrictions using LUA scripts
+     * Processing with Spring Cloud Gateway official script
+     *
+     * @param QPS        queries per second
+     * @param requestUri request uri
+     * @return true : allow access to
+     */
+    @Override
+    public boolean tryAcquire(Double QPS, String requestUri) {
+        try {
+            // get keys
+            List<String> keys = getKeys(requestUri);
+
+            // Parameters and serialization of return values
+            RedisSerializer<String> serializer = new StringRedisSerializer();
+
+            // execute lua script
+            List<Long> tokenResult = (List<Long>) this.redisTemplate.execute(this.redisScript, serializer, serializer, keys, String.valueOf(QPS.intValue()), String.valueOf(QPS.intValue()), Instant.now().getEpochSecond() + "", "1");
+
+            // Index 1 value is the number of remaining requestable tokens
+            Long newTokens = tokenResult.get(1);
+            if (logger.isDebugEnabled()) {
+                logger.debug("Number of remaining tokens for this request is {}", newTokens);
+            }
+            return newTokens > 0;
+        } catch (Exception e) {
+            /*
+             * We don't want a hard dependency on Redis to allow traffic. Make sure to set
+             * an alert so you know if this is happening too much. Stripe's observed
+             * failure rate is 0.01%.
+             */
+            logger.error("Error determining if user allowed from redis", e);
+        }
+        return true;
+    }
+
+    /**
+     * get Redis Script
+     *
+     * @return RedisScript
+     */
+    RedisScript<List<Long>> getRedisScript() {
+        DefaultRedisScript redisScript = new DefaultRedisScript<>();
+        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(QPS_LUA_PATH)));
+        redisScript.setResultType(List.class);
+        return redisScript;
+    }
+
+    /**
+     * get Keys
+     *
+     * @param id resource key（request uri）
+     * @return
+     */
+    static List<String> getKeys(String id) {
+        // use `{}` around keys to use Redis Key hash tags
+        // this allows for using redis cluster
+
+        // Make a unique key per user.
+        String prefix = "qps_rate_limiter.{" + id;
+
+        // You need two Redis keys for Token Bucket.
+        String tokenKey = prefix + "}.tokens";
+        String timestampKey = prefix + "}.timestamp";
+        return Arrays.asList(tokenKey, timestampKey);
+    }
+}