com.github.azbh111.utils.java.lang.RateLimit Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of utils-java Show documentation
com.github.azbh111:utils-java
The newest version!
package com.github.azbh111.utils.java.lang;

import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 限流
 *
 * @author: zyp
 * @since: 2021/11/3 17:07
 */
public class RateLimit {
    private final int seconds;
    private final int limit;
    /**
     * 每隔多久生成1个令牌
     */
    private final int timePerPermit;
    private final int fullTime;
    private AtomicBoolean lock = new AtomicBoolean(false);
    private volatile long lastAcquireTime = 0;
    private volatile long failTimes = 0;

    /**
     * @param seconds seconds 时间窗
     * @param limit   limit 时间窗内可以触发多少次
     * @return null
     * @author zhengyongpan
     * @since 2021/11/3 17:38
     */
    public RateLimit(int seconds, int limit) {
        if (seconds <= 0) {
            throw new RuntimeException("minute must greater than 0");
        }
        if (limit <= 0) {
            throw new RuntimeException("limit must greater than 0");
        }
        this.seconds = seconds;
        this.limit = limit;
        this.fullTime = seconds * 1000;
        this.timePerPermit = seconds * 1000 / limit;
    }

    /**
     * 获取1个令牌
     *
     * @return int <0: 失败. >=0:
     * @author zhengyongpan
     * @since 2021/11/3 17:22
     */
    public long tryGet() {
        return tryGet(1);
    }

    /**
     * 获取指定令牌数量
     * 若获取令牌失败, 返回负数
     * 若获取令牌成功, 返回从上次成功到本次成功这期间失败了多少次
     *
     * @param acquire 获取令牌数量
     * @return int <0: 失败. >=0:
     * @author zhengyongpan
     * @since 2021/11/3 17:22
     */
    public long tryGet(int acquire) {
        if (acquire <= 0) {
            throw new RuntimeException("acquire must greater than 0");
        }
        while (true) {
//            获取锁
            if (!lock.compareAndSet(false, true)) {
//                没有获取到, 让出CPU
                Thread.yield();
                continue;
            }
            try {
                long time = System.currentTimeMillis();
                long lastAcquireTime = this.lastAcquireTime;
                long beginTime = time - fullTime;
                long preAcquireTime = Math.max(lastAcquireTime, beginTime);
                long costTime = acquire * timePerPermit;
                long nextAcquireTime = preAcquireTime + costTime;
                long res;
                if (nextAcquireTime <= time) {
//            令牌足够
                    this.lastAcquireTime = nextAcquireTime;
                    res = this.failTimes;
                    this.failTimes = 0;
                } else {
//            令牌不足
                    this.failTimes++;
                    res = -this.failTimes;
                }
                return res;
            } finally {
                lock.set(false);
            }
        }
    }

}