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

1.令牌桶算法

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

2. 代码实现

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

/**
 * 用于放令牌的桶
 */
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;
    }
}

• RateLimitService 限速的主要方法

/**
 * 限速1s内10个请求
 */
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<String, TokenBucket>();

    //针对某一个用户进行限速
    public boolean available(String key) {
        //1.如果缓存中没有相应key的tokenBucket,需要初始化一个,填充时间为当前系统时间，令牌个数为最大数减一（因为为减去这一次的请求）
        if (!bucketMap.containsKey(key)) {
            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.get(key);
            //获取桶的信息--上次填充距当前时间的间隔，用于计算可以生成多少令牌
            long lastRefillTime = tokenBucket.getLastRefillTime();
            long refillTime = System.currentTimeMillis();
            long intervalSinceLast = refillTime - lastRefillTime;

            long tokenNum;
            //3.如果时间间隔大于1s，也就是要限速的时间粒度，则令牌桶重置
            if (intervalSinceLast > intervalInMills) {
                tokenNum = permits;
            } else { //4.如果时间间隔不大于1s，则需要计算这段时间内生成的令牌个数
                long generatedToken = intervalSinceLast / intervalPerPermit;
                //有可能只取了一个令牌，然后一直生成，之和大于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;
            }
        }
    }
}

• RateLimiteServiceTest

public class RateLimitServiceTest {

    public static RateLimitService rateLimitService = new RateLimitService();

    public static void main(String[] args) {

        while (true) {
            rateLimitService.available("bj");
            try {
                //1.始终有9个,没有超速
                TimeUnit.MILLISECONDS.sleep(100);
                /*
                第一次请求完，桶中还剩9个令牌
                桶中还剩9个令牌
                桶中还剩9个令牌
                桶中还剩9个令牌
                桶中还剩9个令牌
                桶中还剩9个令牌
                桶中还剩9个令牌
                 */
                //2.超速
                //TimeUnit.MILLISECONDS.sleep(10);
                /*
                第一次请求完，桶中还剩9个令牌
                桶中还剩8个令牌
                桶中还剩7个令牌
                桶中还剩6个令牌
                桶中还剩5个令牌
                桶中还剩4个令牌
                桶中还剩3个令牌
                桶中还剩2个令牌
                桶中还剩1个令牌
                桶中还剩0个令牌
                桶中的token已经消费完毕，您已超速。
                 */
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}

3 主要应用场景

• google的RateLimiter有用到令牌桶算法
• SparkStreaming的backpressure机制中同样用到了令牌桶的算法