[开发工具]用JAVA实现一个简单的预热功能

背景

某服务调用，因服务器性能问题，无法直接使用最大qps进行调用，需要动态加速

逻辑

设置最大加速时间，设置允许加速到的最大qps

代码

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.atomic.AtomicLong;

/**
 * 预热控制器
 * @author chunyang.leng
 * @date 2022-03-11 6:22 PM
 */
public class WarmUpManager {
    private static final Logger logger = LoggerFactory.getLogger(WarmUpManager.class);

    /**
     * 预热总时长，单位秒
     */
    private final Long second;

    /**
     * 允许加速到的最大的qps
     */
    private final Long maxQps;

    /**
     * 描述信息
     */
    private String description;

    /**
     * 允许每秒执行的次数
     */
    private final AtomicLong secondQps;
    /**
     * 执行次数计数器
     */
    private AtomicLong requestCount = new AtomicLong(0);

    /**
     * 当前时间,预热功能使用
     */
    private long currentTimeMillis = System.currentTimeMillis();

    /**
     * 时间窗口ms，固定值，一秒
     */
    private final long interval = 1000;

    /**
     * 初始化时间
     */
    private final long initTime = System.currentTimeMillis();

    /**
     * 步进值，默认10
     */
    private final long warmUpLimitStepValue;

    /**
     * 预热结束时间戳
     */
    private final long endTimeMillis;

    /**
     * 创建预热控制对象
     *
     * @param second 预热总时长，单位：秒
     * @param maxQps 允许加速到的最大的qps，不允许为空
     * @param description 描述信息，用于打印日志
     */
    public WarmUpManager(long second,  long maxQps, String description) {
        if(maxQps <= 0){
            throw new IllegalArgumentException("maxQPs 必须大于0");
        }
        this.second = second;
        this.maxQps = maxQps;

        // 计算每秒允许通过的qps
        secondQps = new AtomicLong(maxQps / second);
        if (secondQps.get() == 0) {
            // 防止为0
            secondQps.set(1);
        }
        // 设置步进
        warmUpLimitStepValue = secondQps.get();
        if(description == null){
            this.description = "";
        }
        // 计算预热结束时间
        endTimeMillis = currentTimeMillis + second * interval;
    }


    /**
     * 开始预热
     *
     */
    public void warmUp() throws InterruptedException {
        // 当前时间戳
        long now = System.currentTimeMillis();

        if (now >= endTimeMillis ) {
            // 当前时间大于等于结束时间，预热结束
            return;
        }
        // 计算下一跳窗口起始时间戳
        long total = currentTimeMillis + interval;

        if ((now < total) && requestCount.incrementAndGet() > secondQps.get()) {
            // 当前时间小于下一窗口开始时间，并且请求次数大于限制的qps数,需要进行限速

            // ((now - initTime) / 1000) + 1 ,是因为开始时间为0
            // 日志secondQps.get() + 1 ,是因为 secondQps.getAndAdd(warmUpLimitStepValue); 会丢失1
            LogUtils.info(logger, "{} 预热中,已预热{}秒,qps限流:{}", description, ((now - initTime) / 1000) + 1, secondQps.get() + 1);

            // 执行sleep，睡眠时间为，下一跳开始时间 到 当前时间差
            Thread.sleep(total - now);

            currentTimeMillis = System.currentTimeMillis();
            // 超时后重置
            requestCount = new AtomicLong(0);

            // 重新计算qps大小，步进累加，且不能超出最大限制
            if (secondQps.get() + warmUpLimitStepValue < maxQps) {
                secondQps.getAndAdd(warmUpLimitStepValue);
            }
        }
    }
}

使用方式

public class Test {


    public static void main(String[] args) throws InterruptedException {
        WarmUpManager warmUpManager = new WarmUpManager(10L,5L,"");

        int max = 1000;
        for (int i = 0; i < max; i++) {
            warmUpManager.warmUp();
            System.out.println("==========> " + (  i));
        }
        System.out.println("====run done ");
    }

}