Redis令牌桶算法在限速中的应用

1.令牌桶算法

令牌桶中有初始容量，每来一个请求从桶中获取一个令牌，并且在一定时间间隔中可以生成令牌，多余的令牌被丢弃。可以实现限速功能。

 

2. 使用google的guava缓存，设置缓存失效时间

可参考：https://www.imooc.com/article/28865

 

3. 代码实现

主要针对不同用户的请求进行限速，如果单独使用google的RateLimiter可以控制请求的速率，如果超过限定的速率则进行等待，但是无法获取用户的请求速率。如果下可以根据不同的用户进行限速。

public class TokenBucket {

    //上次填充的时间戳,ms
    private Long lastRefillTime;
    //桶中剩余的token数量
    private Long remainingToken;

    public TokenBucket(long lastRefillTime,long remainingToken){
        this.lastRefillTime = lastRefillTime;
        this.remainingToken = remainingToken;
    }

    public Long getLastRefillTime() {
        return lastRefillTime;
    }

    public Long getRemainingToken() {
        return remainingToken;
    }

    public void setLastRefillTime(Long lastRefillTime) {
        this.lastRefillTime = lastRefillTime;
    }

    public void setRemainingToken(Long remainingToken) {
        this.remainingToken = remainingToken;
    }
}

 

限流主要算法：

public class RateLimitService {
    //每秒的最大次数
    private final long permits = 10;
    //1s的时间间隔
    private final long intervalInMills = 1 * 1000;
    //令牌的生成速率 100ms
    private final long intervalPerPermit = intervalInMills / permits;

    //用作缓存
//    public static ConcurrentMap<String, TokenBucket> bucketMap = new ConcurrentHashMap<>();
    private static Cache<String,TokenBucket> bucketMap = CacheBuilder.newBuilder().initialCapacity(10).concurrencyLevel(1).expireAfterWrite(10, TimeUnit.MINUTES).build();

    //针对某一个用户进行限速
    public boolean available(String key) {
        //1.如果缓存中没有相应key的tokenBucket,需要初始化一个,填充时间为当前系统时间，令牌个数为最大数减一（因为为减去这一次的请求）
        if (bucketMap.getIfPresent(key) == null) {
            TokenBucket tokenBucket = new TokenBucket(System.currentTimeMillis(), permits - 1);
            bucketMap.put(key, tokenBucket);
            System.out.println(String.format("第一次请求完，桶中还剩%s个令牌", tokenBucket.getRemainingToken()));
            return true;
        } else {  //2.缓存中存在
            TokenBucket tokenBucket = bucketMap.getIfPresent(key);
            //获取桶的信息--上次填充距当前时间的间隔，用于计算可以生成多少令牌
            long lastRefillTime = tokenBucket.getLastRefillTime();
            long refillTime = System.currentTimeMillis();
            long intervalSinceLast = refillTime - lastRefillTime;
            System.out.println(String.format("时间差%s",intervalSinceLast));

            long tokenNum;
            //3.如果时间间隔大于1s，也就是要限速的时间粒度，则令牌桶重置
            if (intervalSinceLast > intervalInMills) {
                tokenNum = permits;
            } else { //4.如果时间间隔小于1s，则需要计算这段时间内生成的令牌个数
                long generatedToken = intervalSinceLast / intervalPerPermit;
                System.out.println(String.format("token生成%s个",generatedToken));
                //有可能只取了一个令牌，然后一直生成，之和大于bucket的最大值
                tokenNum = Math.min(permits, tokenBucket.getRemainingToken() + generatedToken);
            }

            //设置令牌桶中的remainingToken和填充时间
            tokenBucket.setLastRefillTime(refillTime);
            //桶中没有token了，并且上次获取token完距现在不足以生成一个token
            if (tokenNum == 0) {
                tokenBucket.setRemainingToken(tokenNum);
                System.out.println("桶中的token已经消费完毕，您已超速。");
                bucketMap.put(key,tokenBucket);
                return false;
            } else {
                tokenBucket.setRemainingToken(tokenNum - 1);
                System.out.println(String.format("桶中还剩%s个令牌", tokenBucket.getRemainingToken()));
                bucketMap.put(key,tokenBucket);
                return true;
            }
        }
    }

}

 

测试验证

public class RateLimitServiceTest {
    public static RateLimitService rateLimitService = new RateLimitService();

    public static void main(String[] args) {
        while (true) {
            rateLimitService.available("R001");
            try {
                /**
                 * 100
                 * 第一次请求完，桶中还剩9个令牌
                 * 第一次请求完，桶中还剩9个令牌
                 * 第一次请求完，桶中还剩9个令牌
                 * 10
                 * 第一次请求完，桶中还剩9个令牌
                 * 时间差20
                 * token生成0个
                 * 桶中还剩8个令牌
                 * 时间差6
                 * token生成0个
                 * 桶中还剩7个令牌
                 * 时间差3
                 * token生成0个
                 * 桶中还剩6个令牌
                 * 时间差2
                 * token生成0个
                 * 桶中还剩5个令牌
                 * 时间差2
                 * token生成0个
                 * 桶中还剩4个令牌
                 * 时间差2
                 * token生成0个
                 * 桶中还剩3个令牌
                 * 时间差2
                 * token生成0个
                 * 桶中还剩2个令牌
                 * 时间差2
                 * token生成0个
                 * 桶中还剩1个令牌
                 * 时间差2
                 * token生成0个
                 * 桶中还剩0个令牌
                 * 时间差2
                 */
                TimeUnit.MICROSECONDS.sleep(10);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}

 

 

 

 

3 主要应用场景

• google的RateLimiter有用到令牌桶算法
• SparkStreaming的backpressure机制中同样用到了令牌桶的算法
• 网关中，相同sessionId每秒限流