漏斗算法限流Java实现

/**
 * 漏斗算法限流
 * @author sixiaojie
 * @date 2021-08-11-11:00
 */
@Slf4j
@Component
public class FunnelRateLimiter {
    static class Funnel{
        /** 漏斗容量 */
        int capacity;
        /** 漏嘴流水速率 */
        float leakingRate;
        /** 漏斗剩余空间 */
        int leftQuota;
        /** 上一次漏水时间 */
        long leakingTs;

        public Funnel(int capacity, float leakingRate) {
            this.capacity = capacity;
            this.leakingRate = leakingRate;
            this.leftQuota = capacity;
            this.leakingTs = System.currentTimeMillis();
        }
        void makeSpace(){
            long nowTs = System.currentTimeMillis();
            // 距离上一次漏水过去了多久
            long deltaTs = nowTs - leakingTs;
            int deltaQuota = (int) (deltaTs * leakingRate);
            // 间隔时间过长，整数数字过大溢出
            if(deltaQuota < 0){
                this.leftQuota = capacity;
                this.leakingTs = nowTs;
                return;
            }
            // 腾出空间太小，最小单位是1，那就等下次吧
            if(deltaQuota < 1){
                return;
            }
            // 增加剩余空间
            this.leftQuota += deltaQuota;
            // 记录漏水时间
            this.leakingTs = nowTs;
            // 剩余空间不得高于容量
            if(this.leftQuota > this.capacity){
                this.leftQuota = this.capacity;
            }
        }

        boolean watering(int quota){
            makeSpace();
            // 判断剩余空间是否足够
            if(this.leftQuota >= quota){
                this.leftQuota -= quota;
                return true;
            }
            return false;
        }
    }

    /**
     * 所有的漏斗
     */
    private Map<String,Funnel> funnels = new HashMap<>();

    /**
     *
     * @param userId 用户ID
     * @param actionKey 行为标示
     * @param capacity 漏斗容量
     * @param leakingRate 漏嘴流水速率 quota / s
     * @return
     */
    public boolean isActionAllowed(String userId,String actionKey,int capacity,float leakingRate){
        String key = String.format("%s:%s", userId, actionKey);
        Funnel funnel = funnels.get(key);
        if(funnel == null){
            funnel = new Funnel(capacity, leakingRate/1000);
            funnels.put(key,funnel);
        }
        // 需要1个quota
        return funnel.watering(1);
    }

    public static void main(String[] args) throws InterruptedException {
        FunnelRateLimiter funnelRateLimiter = new FunnelRateLimiter();
        while (true){
            // 漏嘴流水速率 每分钟10个
            boolean actionAllowed = funnelRateLimiter.isActionAllowed("sixj", "sms", 10, 10 / 60.0f);
            log.info("--->"+actionAllowed);
            Thread.sleep(1000);
        }
    }
}