基于redis的分布式令牌桶

令牌桶
单机：

private RateLimiter rateLimiter = RateLimiter.create(300L);不适合集群环境

分布式：

package com.taoheng.util;
import lombok.extern.slf4j.Slf4j;

/**
 * 基于redis的分布式令牌桶
 * @Description
 * @Author taoheng
 * @Date 2019/7/26 16:55
 **/

@Slf4j
public class RedisRateLimiter {

    /**
     * 请大家把自己定义的“令牌桶”统一放到这里。
     *
     * 解释：为了防止多个业务功能命名的“令牌桶名称”相同，而导致共用一个令牌桶。
     * 或同一个令牌桶用了不同的名称，而导致不能实现真正的限流作用。
     */
    public enum TokenBucketEnum{
        /** 云鸽消息API接口限流（300次/秒） */
        MSC("msc", 300L, 1, 10L);

        private String BASE_NAME = "yunzhi_aecp_";
        /** 令牌桶名称 */
        private String name;
        /** 令牌数 */
        private Long permits;
        /** 周期（秒）*/
        private Integer cycle;
        /** 自旋阻塞时间(毫秒)，建议为周期的1/100 */
        private Long spinBlockingTime;
        TokenBucketEnum (String name, Long permits, Integer cycle, Long spinBlockingTime){
            this.name = BASE_NAME+name;
            this.permits = permits;
            this.cycle = cycle;
            this.spinBlockingTime = spinBlockingTime;
        }
    }

    /**
     * 阻塞式获取令牌
     * @param tokenBucket 令牌桶
     * @throws Exception
     */
    public static void acquire(TokenBucketEnum tokenBucket) throws Exception {
        Redis redis = SpringUtil.getBean("redis", Redis.class);
        /*
         * 每次对key进行递增，如果递增结果大于规定的QPS就进行自旋阻塞
         *
         * 在自旋途中key可能会因超过1s而失效，而根据incr命令的特性（如果 key 不存在，那么 key 的值会先被初始化为 0 ，然后再执行 INCR 操作），
         * 这时会有一次自旋操作返回1，于是跳出阻塞。
         *
         * 由第一个获取值为1的客户端实例来将key设置过期时间为1s
         */
        //每次自增操作后的值
        Long incrementalResult ;
        //自旋次数
        int spinNum = 0;
        while ((incrementalResult = redis.incr(tokenBucket.name)) > tokenBucket.permits){
            /*
             * 如果本次自旋时间大于令牌桶的一个周期,说明上一次的key设置失效时间失败，导致key不会自动失效，进而导致这里出现死循环
             * 此时需要重置令牌桶
             */
            if(spinNum * tokenBucket.spinBlockingTime > tokenBucket.cycle * 1000){
                redis.del(tokenBucket.name);
                //重新计数
                spinNum = 0;
                log.warn("{}令牌桶出现死循环，已被重置", tokenBucket.name);
                continue;
            }
            //自旋阻塞
            Thread.sleep(tokenBucket.spinBlockingTime);
            spinNum++;
        }
        /*
         * 由于incr()和expire()是分步操作（非原子性），所以在这两步操作的时间间隙内可能会有其它客户端对此令牌桶进行了自增操作。
         * 极端情况下会在expire()操作之前达到令牌桶上限，但你并不需要担心这个问题，因为达到令牌桶上限后它们都会自动进入自旋阻塞，而不会突破令牌桶的限制。
         * 等待令牌桶的再次因失效后被重置后，它们就又可以再次获取令牌了。
         * 综上：次令牌桶的实际限流是=令牌数/（一个周期+两步操作的时间间隙），即不会超过你设置的阀值。
         */
        if(incrementalResult == 1){
            //设置“tokenBucketCycle”后失效
            redis.expire(tokenBucket.name, tokenBucket.cycle);
        }
    }
}