Java基于时间窗口的限流算法

Java基于时间窗口的限流算法

算法流程

CurrentLimiterInputStream

package com.liangzhm.io;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;

/**
 * @Classname CurrentLimiterInputStream
 * @Description TODO
 * @Date 2023/1/18 10:22
 * @Author liangzhm
 * @Version 1.0
 */
public class CurrentLimiterInputStream extends BufferedInputStream {

    private static int DEFAULT_BUFFER_SIZE = 8192;

    private CurrentLimiterVo streamLimitVo;

    public CurrentLimiterInputStream(InputStream inputStream, CurrentLimiterVo streamLimitVo) {
        this(inputStream, DEFAULT_BUFFER_SIZE, streamLimitVo);
    }

    public CurrentLimiterInputStream(InputStream inputStream, int size) {
        this(inputStream, size, null);
    }

    public CurrentLimiterInputStream(InputStream inputStream, int size, CurrentLimiterVo streamLimitVo) {
        super(inputStream, size);
        this.streamLimitVo = streamLimitVo;
    }

    @Override
    public int read(byte b[], int off, int len) throws IOException {
        int bytes = super.read(b, off, len);
        if (streamLimitVo != null) {
            streamLimitVo.limit(bytes);
        }
        return bytes;
    }
}

CurrentLimiterIoUtils

package com.liangzhm.io;

import org.springframework.util.StreamUtils;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

/**
 * @Classname CurrentLimiterIoUtils
 * @Description TODO
 * @Date 2023/1/18 13:47
 * @Author liangzhm
 * @Version 1.0
 */
public class CurrentLimiterIoUtils {

    /**
     * @param in
     * @param out
     * @param maxRate 最大速率，单位kb,如：限速200kb/s，则maxRate填写200
     * @throws IOException
     */
    public static void copy(InputStream in, OutputStream out, int maxRate) throws IOException {
        copy(in, out, CurrentLimiterVo.of(maxRate));
    }


    /**
     * @param in
     * @param out
     * @param maxRate    最大速率，单位kb,如：限速200kb/s，则maxRate填写200
     * @param timeWindow 时间窗口，单位秒,如：限速200kb/s，则timeWindow填写1
     * @throws IOException
     */
    public static void copy(InputStream in, OutputStream out, int maxRate, int timeWindow) throws IOException {
        copy(in, out, CurrentLimiterVo.of(timeWindow, maxRate));
    }

    /**
     * @param in
     * @param out
     * @param currentLimiterVo
     * @throws IOException
     */
    public static void copy(InputStream in, OutputStream out, CurrentLimiterVo currentLimiterVo) throws IOException {
        StreamUtils.copy(new CurrentLimiterInputStream(in, currentLimiterVo), out);
    }
}

CurrentLimiterVo

package com.liangzhm.io;

import java.util.concurrent.TimeUnit;

/**
 * @Classname CurrentLimiterVo
 * @Description TODO
 * @Date 2023/1/18 12:07
 * @Author liangzhm
 * @Version 1.0
 */
public class CurrentLimiterVo {

    private static final int ONE = 1;

    private static final int KB = 1024;

    private static final int MILLION = 1000_000;

    /**
     * 纳秒
     */
    private long timeWindow;

    private int chunk;

    private long previousTime;

    private int currentBytes = 0;

    /**
     * 单位KB
     */
    private int maxRate;

    private CurrentLimiterVo(int maxRate) {
        this(ONE, maxRate);
    }

    private CurrentLimiterVo(int timeWindow, int maxRate) {
        if (timeWindow < 0) {
            throw new RuntimeException("timeWindow不合法");
        }
        if (maxRate < 0) {
            throw new RuntimeException("maxRate不合法");
        }
        this.maxRate = maxRate;
        this.timeWindow = TimeUnit.SECONDS.toNanos(timeWindow);
        this.chunk = this.maxRate * KB;
    }

    /**
     * @param maxRate
     * @return
     */
    public static CurrentLimiterVo of(int maxRate) {
        return new CurrentLimiterVo(maxRate);
    }

    /**
     * @param timeWindow
     * @param maxRate
     * @return
     */
    public static CurrentLimiterVo of(int timeWindow, int maxRate) {
        return new CurrentLimiterVo(timeWindow, maxRate);
    }

    public void limit(int bytes) {
        if (bytes <= 0) {
            return;
        }
        this.currentBytes += bytes;
        if (this.previousTime == 0) {
            this.previousTime = System.nanoTime();
        }
        while (this.currentBytes >= chunk) {
            long passTime = System.nanoTime() - this.previousTime;
            long missedTime = this.timeWindow - passTime;
            if (missedTime > 0) {
                try {
                    Thread.sleep(missedTime / MILLION, (int) (missedTime % MILLION));
                } catch (InterruptedException e) {
                }
            }
            this.currentBytes -= chunk;
            this.previousTime = System.nanoTime();
        }
    }
}