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每秒限流