提交代码

build.gradle

@@ -38,7 +38,7 @@ allprojects {
 	}
 
 	group = "cn.com.duiba.boot"
-	version = "0.0.74-sjq-SNAPSHOT"
+	version = "0.0.80-sjq-SNAPSHOT"
 }
 
 subprojects {

spring-boot-starter-dsp-rateLimiter/src/main/java/cn/com/duiba/spring/boot/starter/dsp/rateLimiter/Commons.java

@@ -18,4 +18,9 @@ public class Commons {
         return redisScript;
     }
 
+    @Bean
+    public TokenBucketLimiter createTokenBucketLimiter() {
+        return new TokenBucketLimiter();
+    }
+
 }

spring-boot-starter-dsp-rateLimiter/src/main/java/cn/com/duiba/spring/boot/starter/dsp/rateLimiter/LimitAspect.java

@@ -8,13 +8,8 @@ import org.aspectj.lang.reflect.MethodSignature;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.data.redis.core.StringRedisTemplate;
-import org.springframework.data.redis.core.script.DefaultRedisScript;
 import org.springframework.stereotype.Component;
-import javax.annotation.Resource;
 import java.lang.reflect.Method;
-import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
@@ -30,7 +25,6 @@ public class LimitAspect {
     private final static Map<String, AtomicLong> map = new ConcurrentHashMap<>();
 
     private final static Integer REDIS_NODE_NUM = 4;
-
     static {
         map.put("rate.limit:com.duiba.tuia.adx.web.service.algo.impl.AdxAlgoServiceImpl-hello-limit", new AtomicLong());
 
@@ -52,14 +46,11 @@ public class LimitAspect {
         redisNodeIndex.put(15L, 47);
     }
 
-    @Resource(name = "redis03StringRedisTemplate")
-    private StringRedisTemplate stringRedisTemplate;
-
     @Autowired
-    private DefaultRedisScript<Long> redisLuaScript;
+    private RateLimitProperties rateLimitProperties;
 
     @Autowired
-    private RateLimitProperties rateLimitProperties;
+    private TokenBucketLimiter tokenBucketLimiter;
 
     @Pointcut(value = "@annotation(cn.com.duiba.spring.boot.starter.dsp.rateLimiter.RateLimit)")
     public void rateLimitPointcut() {
@@ -94,15 +85,13 @@ public class LimitAspect {
         String redisKey = commonRedisKey + "{" + redisNodeIndex.get(index) + "}";
         logger.info("限流啦, redis key{}", redisKey);
 
-        List<String> keys = Collections.singletonList(redisKey);
-
-        int totalLimitCount = rateLimit.count();
-        int limitCount = totalLimitCount % REDIS_NODE_NUM > index ? totalLimitCount / REDIS_NODE_NUM + 1 : totalLimitCount / REDIS_NODE_NUM;
-
-        Long number = stringRedisTemplate.execute(redisLuaScript, keys, String.valueOf(limitCount), String.valueOf(rateLimit.time()));
-
-        if (number != null && number != 0 && number <= limitCount) {
-            logger.info("限流时间段内访问第：{} 次", number);
+        long totalLimitCount = rateLimit.count();
+        long limitCount = totalLimitCount % REDIS_NODE_NUM > index ? totalLimitCount / REDIS_NODE_NUM + 1 : totalLimitCount / REDIS_NODE_NUM;
+        if (limitCount <= 0) {
+            throw new RuntimeException("已经到设置限流次数");
+        }
+        TokenBucketLimiterPolicy tokenBucketLimiterPolicy = new TokenBucketLimiterPolicy(limitCount, 200);
+        if (tokenBucketLimiter.acquire(redisKey, tokenBucketLimiterPolicy)) {
             return joinPoint.proceed();
         }
 

spring-boot-starter-dsp-rateLimiter/src/main/java/cn/com/duiba/spring/boot/starter/dsp/rateLimiter/RateLimit.java

@@ -11,8 +11,8 @@ public @interface RateLimit {
 
     String key() default "limit";
 
-    int time() default 5;
+    long time() default 5;
 
-    int count() default 5;
+    long count() default 5;
 
 }

spring-boot-starter-dsp-rateLimiter/src/main/java/cn/com/duiba/spring/boot/starter/dsp/rateLimiter/TokenBucketLimiter.java

+package cn.com.duiba.spring.boot.starter.dsp.rateLimiter;
+
+import com.google.common.collect.Lists;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.redis.core.StringRedisTemplate;
+import org.springframework.data.redis.core.script.DefaultRedisScript;
+import javax.annotation.Resource;
+import java.util.Objects;
+
+@Slf4j
+public class TokenBucketLimiter {
+
+    @Resource(name = "redis03StringRedisTemplate")
+    private StringRedisTemplate stringRedisTemplate;
+
+    @Autowired
+    private DefaultRedisScript<Long> redisLuaScript;
+
+    boolean acquire(String key, TokenBucketLimiterPolicy policy) {
+        if (policy == null) {
+            return true;
+        }
+        Long remain = stringRedisTemplate.execute(redisLuaScript, Lists.asList(key, new String[]{}), policy.toParams());
+        return Objects.nonNull(remain) && remain > 0;
+    }
+
+}

spring-boot-starter-dsp-rateLimiter/src/main/java/cn/com/duiba/spring/boot/starter/dsp/rateLimiter/TokenBucketLimiterPolicy.java

+package cn.com.duiba.spring.boot.starter.dsp.rateLimiter;
+
+import com.google.common.collect.Lists;
+import lombok.Data;
+import java.util.List;
+
+/**
+ * 令牌桶限流器的执行对象
+ */
+@Data
+public class TokenBucketLimiterPolicy {
+
+    /**
+     * 限流时间间隔
+     * (重置桶内令牌的时间间隔)
+     */
+    private long resetBucketInterval;
+    /**
+     * 最大令牌数量
+     */
+    private long bucketMaxTokens;
+
+    /**
+     * 初始可存储数量
+     */
+    private long initTokens;
+
+    /**
+     * 每个令牌产生的时间
+     */
+    private long intervalPerPermit;
+
+    public TokenBucketLimiterPolicy(long bucketMaxTokens, long maxBurstTime) {
+        // 最大令牌数
+        this.bucketMaxTokens = bucketMaxTokens;
+        // 限流时间间隔
+        this.resetBucketInterval = 1000;
+        if (bucketMaxTokens >= 200) {
+            intervalPerPermit = 50;
+            initTokens = bucketMaxTokens / 50;
+            return;
+        }
+        // 令牌的产生间隔 = 限流时间 / 最大令牌数
+        intervalPerPermit = resetBucketInterval / bucketMaxTokens;
+        // 初始令牌数 = 最大的突发流量的持续时间 / 令牌产生间隔
+        // 用 最大的突发流量的持续时间 计算的结果更加合理,并不是每次初始化都要将桶装满
+        initTokens = Math.min(Math.max(maxBurstTime * bucketMaxTokens / resetBucketInterval, 1), bucketMaxTokens);
+    }
+
+    public String[] toParams() {
+        List<String > list = Lists.newArrayList();
+        list.add(String.valueOf(getIntervalPerPermit()));
+        list.add(String.valueOf(System.currentTimeMillis()));
+        list.add(String.valueOf(getInitTokens()));
+        list.add(String.valueOf(getBucketMaxTokens()));
+        list.add(String.valueOf(getResetBucketInterval()));
+        return list.toArray(new String[]{});
+    }
+
+}

spring-boot-starter-dsp-rateLimiter/src/main/resources/limit.lua

-local limit = tonumber(ARGV[1])        --限流大小
-local current = tonumber(redis.call('get', KEYS[1]) or "0")
-if current + 1 > limit then --如果超出限流大小
-  return 0
-else  --请求数+1，并设置2秒过期
-  redis.call("INCRBY", KEYS[1], "1")
-   redis.call("expire", KEYS[1], "2")
-   return current + 1
-end
+local intervalPerTokens = tonumber(ARGV[1])
+local curTime = tonumber(ARGV[2])
+local initTokens = tonumber(ARGV[3])
+local bucketMaxTokens = tonumber(ARGV[4])
+local resetBucketInterval = tonumber(ARGV[5])
+
+local bucket = redis.call('hgetall', KEYS[1])
+local currentTokens
+
+-- 若当前桶未初始化,先初始化令牌桶
+if table.maxn(bucket) == 0 then
+    -- 初始桶内令牌
+    currentTokens = initTokens
+    -- 设置桶最近的填充时间是当前
+    redis.call('hset', KEYS[1], 'lastRefillTime', curTime)
+    -- 初始化令牌桶的过期时间, 设置为间隔的 1.5 倍
+    redis.call('pexpire', KEYS[1], resetBucketInterval * 1.5)
+
+-- 若桶已初始化,开始计算桶内令牌
+-- 为什么等于 4 ? 因为有两对 field, 加起来长度是 4
+-- { "lastRefillTime(上一次更新时间)","curTime(更新时间值)","tokensRemaining(当前保留的令牌)","令牌数" }
+elseif table.maxn(bucket) == 4 then
+
+    -- 上次填充时间
+    local lastRefillTime = tonumber(bucket[2])
+    -- 剩余的令牌数
+    local tokensRemaining = tonumber(bucket[4])
+
+    -- 当前时间大于上次填充时间
+    if curTime > lastRefillTime then
+
+        -- 拿到当前时间与上次填充时间的时间间隔
+        -- 举例理解: curTime = 2620 , lastRefillTime = 2000, intervalSinceLast = 620
+        local intervalSinceLast = curTime - lastRefillTime
 
+        -- 如果当前时间间隔 大于 令牌的生成间隔
+        -- 举例理解: intervalSinceLast = 620, resetBucketInterval = 1000
+        if intervalSinceLast > resetBucketInterval then
 
+            -- 将当前令牌填充满
+            currentTokens = initTokens
 
+            -- 更新重新填充时间
+            redis.call('hset', KEYS[1], 'lastRefillTime', curTime)
 
+        -- 如果当前时间间隔 小于 令牌的生成间隔
+        else
 
+            -- 可授予的令牌 = 向下取整数( 上次填充时间与当前时间的时间间隔 / 两个令牌许可之间的时间间隔 )
+            -- 举例理解 : intervalPerTokens = 200 ms , 令牌间隔时间为 200ms
+            --           intervalSinceLast = 620 ms , 当前距离上一个填充时间差为 620ms
+            --           grantedTokens = 620/200 = 3.1 = 3
+            local grantedTokens = math.floor(intervalSinceLast / intervalPerTokens)
 
---
---
--- -- LUA脚本会以单线程执行,不会有并发问题，一个脚本中的执行过程中如果报错，那么已执行的操作不会回滚
--- -- KEYS和ARGV是外部传入进来需要操作的redis数据库中的key,下标从1开始
--- -- 参数结构: KEYS = [限流的key]   ARGV = [最大令牌数, 每秒生成的令牌数, 本次请求的毫秒数]
--- local info = redis.pcall('HMGET', KEYS[1], 'last_time', 'stored_token_nums')
--- local last_time = info[1] --最后一次通过限流的时间
--- local stored_token_nums = tonumber(info[2]) -- 剩余的令牌数量
--- local max_token = tonumber(ARGV[1])
--- local token_rate = tonumber(ARGV[2])
--- local current_time = tonumber(ARGV[3])
--- local past_time = 0
--- local rateOfperMills = token_rate/1000 -- 每毫秒生产令牌速率
---
--- if stored_token_nums == nil then
---     -- 第一次请求或者键已经过期
---     stored_token_nums = max_token --令牌恢复至最大数量
---     last_time = current_time --记录请求时间
--- else
---     -- 处于流量中
---     past_time = current_time - last_time --经过了多少时间
---
---     if past_time <= 0 then
---         --高并发下每个服务的时间可能不一致
---         past_time = 0 -- 强制变成0 此处可能会出现少量误差
---     end
---     -- 两次请求期间内应该生成多少个token
---     local generated_nums = math.floor(past_time * rateOfperMills)  -- 向下取整，多余的认为还没生成完
---     stored_token_nums = math.min((stored_token_nums + generated_nums), max_token) -- 合并所有的令牌后不能超过设定的最大令牌数
--- end
---
--- local returnVal = 0 -- 返回值
---
--- if stored_token_nums > 0 then
---     returnVal = 1 -- 通过限流
---     stored_token_nums = stored_token_nums - 1 -- 减少令牌
---     -- 必须要在获得令牌后才能重新记录时间。举例: 当每隔2ms请求一次时,只要第一次没有获取到token,那么后续会无法生产token,永远只过去了2ms
---     last_time = last_time + past_time
--- end
---
--- -- 更新缓存
--- redis.call('HMSET', KEYS[1], 'last_time', last_time, 'stored_token_nums', stored_token_nums)
--- -- 设置超时时间
--- -- 令牌桶满额的时间(超时时间)(ms) = 空缺的令牌数 * 生成一枚令牌所需要的毫秒数(1 / 每毫秒生产令牌速率)
--- redis.call('PEXPIRE', KEYS[1], math.ceil((1/rateOfperMills) * (max_token - stored_token_nums)))
---
--- return returnVal
---
--- --实现平滑限流
--- --出自：https://www.jianshu.com/p/89822f8d5c69
--- --https://blog.csdn.net/teavamc/article/details/113359632
\ No newline at end of file
+            -- 可授予的令牌 > 0 时
+            -- 举例理解 : grantedTokens = 620/200 = 3.1 = 3
+            if grantedTokens > 0 then
+
+                -- 生成的令牌 = 上次填充时间与当前时间的时间间隔 % 两个令牌许可之间的时间间隔
+                -- 举例理解 : padMillis = 620%200 = 20
+                --           curTime = 2620
+                --           curTime - padMillis = 2600
+                local padMillis = math.fmod(intervalSinceLast, intervalPerTokens)
+
+                -- 将当前令牌桶更新到上一次生成时间
+                redis.call('hset', KEYS[1], 'lastRefillTime', curTime - padMillis)
+            end
+
+            -- 更新当前令牌桶中的令牌数
+            -- Math.min(根据时间生成的令牌数 + 剩下的令牌数, 桶的限制) => 超出桶最大令牌的就丢弃
+            currentTokens = math.min(grantedTokens + tokensRemaining, bucketMaxTokens)
+        end
+    else
+        -- 如果当前时间小于或等于上次更新的时间, 说明刚刚初始化, 当前令牌数量等于桶内令牌数
+        -- 不需要重新填充
+        currentTokens = tokensRemaining
+    end
+end
+
+-- 如果当前桶内令牌小于 0,抛出异常
+assert(currentTokens >= 0)
+
+-- 如果当前令牌 == 0 ,更新桶内令牌, 返回 0
+if currentTokens == 0 then
+    redis.call('hset', KEYS[1], 'tokensRemaining', currentTokens)
+    return 0
+else
+    -- 如果当前令牌 大于 0, 更新当前桶内的令牌 -1 , 再返回当前桶内令牌数
+    redis.call('hset', KEYS[1], 'tokensRemaining', currentTokens - 1)
+    return currentTokens
+end