令牌桶限流-java实现

令牌桶的容量是c（个），令牌以速度r（个/秒）均匀的放入桐中，上个请求的时间为at（时间戳），上个请求后剩余的令牌数目为w（个），现在有个请求b对象进来了，现在请求的时间bt=now()，伪代码如上图，其中wb代表从at到bt时间段内产生的令牌数，产生的令牌数加上上次剩余的令牌数是不能大于桶容量的
Java实现代码：

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class BucketLimiter {
    /* 桶容量 */
    private final long c;
    /* 令牌放入速度 */
    private final int r;
    /* 桶初始化时间 */
    private final long initTime;

    /* 模拟redis缓存 */
    private final Map<String, Object> redisCache = new HashMap<>(2048);
    /* 请求时间redis key */
    public static final String REQUEST_TIME = "REQUEST_TIME";
    /* 请求后剩余令牌数redis key */
    public static final String TOKEN_REMAINS = "TOKEN_REMAINS";

    /**
     * @param c 桶容量
     * @param r 令牌放入速度
     */
    public BucketLimiter(long c, int r) {
        this.c = c;
        this.r = r;
        this.initTime = System.currentTimeMillis();
    }

    /* 注意这里要加同步 */
    public synchronized boolean permit() {
        //获取上个请求后令牌桶状态
        BucketStatus status = getLastRequestBucketStatus();

        //上个请求时间
        final long at = status.getRequestTime();
        //上个请求后剩余令牌数
        long w = status.getRemains();
        //现请求时间
        final long bt = System.currentTimeMillis();
        //从上个请求到现在请求增加的令牌数
        final long wb = (bt - at)/1000 * r;
        System.out.println("===生产数："+wb);
        //现在桶里面剩余的令牌数
        w = Math.min(w + wb, c);
        //假设每次消耗一个令牌
        if (w > 0) {
            w--;
            //请求时间和剩余令牌数
            redisCache.put(REQUEST_TIME, bt);
            redisCache.put(TOKEN_REMAINS, w);
            return true;
        } else {
            return false;
        }
    }

    /**
     * @description 获取上个请求后令牌桶状态
     * @author yulin.chen
     * @date 2021/10/26 10:27
     */
    private BucketStatus getLastRequestBucketStatus() {
        Object requestTime = redisCache.get(REQUEST_TIME);
        Object remains = redisCache.get(TOKEN_REMAINS);
        if (null == requestTime) {
            //请求时间为空就是第一次请求，剩余数设置为0，时间设置为初始化时间
            return new BucketStatus(0L, this.initTime);
        }
        return new BucketStatus(Long.parseLong(String.valueOf(remains)), Long.parseLong(String.valueOf(requestTime)));
    }

    /**
     *
     */
    public static class BucketStatus {
        /* 剩余令牌数 */
        private final Long remains;
        /* 请求时间 */
        private final Long requestTime;

        public BucketStatus(Long remains, Long requestTime) {
            this.remains = remains;
            this.requestTime = requestTime;
        }

        public Long getRemains() {
            return remains;
        }

        public Long getRequestTime() {
            return requestTime;
        }
    }
}

测试代码

public static void main(String[] args) throws InterruptedException {
    BucketLimiter limiter = new BucketLimiter(5000, 5);
    //生产5秒，每秒生产5个
    TimeUnit.SECONDS.sleep(5);
    for (int i = 0; i < 10000; i++) {
        new Thread(() -> System.out.println(Thread.currentThread().getName() + " request permit:" + limiter.permit()), "[Thread " + i + "]")
                .start();
    }
}