Java实现文件下载/上传的限流算法

在做文件下载功能时，为了避免下载功能将服务器的带宽打满，从而影响服务器的其他服务。我们可以设计一个限流器来限制下载的速率，从而限制下载服务所占用的带宽。

一、算法思路

定义一个数据块 chunk (单位 bytes) 以及允许的最大速率 maxRate (单位 KB/s)。通过 maxRate 我们可以算出，在 maxRate 的速率下，通过一个数据块大小的字节流所需要的时间 timeCostPerChunk。之后，在读取 / 写入字节时，我们维护已经读取 / 写入的字节量 bytesWillBeSentOrReceive。当 bytesWillBeSentOrReceive 达到一个数据块的大小时，检查期间消耗的时间 (nowNanoTime-lastPieceSentOrReceiveTick)。如果期间消耗的时间小于 timeCostPerChunk 的值，说明当前的速率已经超过了 maxRate 的速率，这时候就需要休眠一会来限制流量。如果速率没超过或者休眠完后，将 bytesWillBeSentOrReceive=bytesWillBeSentOrReceive-chunkSize之后在读取 / 写入数据时继续检查。

二、限流的完整 java 代码实现

限流器的实现

package com.example.demo.limit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class BandwidthLimiter {
    Logger LOGGER = LoggerFactory.getLogger(BandwidthLimiter.class);
    //KB代表的字节数
    private static final Long KB = 1024L;
    //一个chunk的大小，单位byte。设置一个块的大小为1M
    private static final Long CHUNK_LENGTH = 1024 * 1024L;

    //已经发送/读取的字节数
    private int bytesWillBeSentOrReceive = 0;
    //上一次接收到字节流的时间戳——单位纳秒
    private long lastPieceSentOrReceiveTick = System.nanoTime();
    //允许的最大速率，默认为 1024KB/s
    private int maxRate = 1024;
    //在maxRate的速率下，通过chunk大小的字节流要多少时间（纳秒）
    private long timeCostPerChunk = (1000000000L * CHUNK_LENGTH) / (this.maxRate * KB);

    public BandwidthLimiter(int maxRate) {
        this.setMaxRate(maxRate);
    }

    //动态调整最大速率
    public void setMaxRate(int maxRate) {
        if (maxRate < 0) {
            throw new IllegalArgumentException("maxRate can not less than 0");
        }
        this.maxRate = maxRate;
        if (maxRate == 0) {
            this.timeCostPerChunk = 0;
        } else {
            this.timeCostPerChunk = (1000000000L * CHUNK_LENGTH) / (this.maxRate * KB);
        }
    }

    public synchronized void limitNextBytes() {
        this.limitNextBytes(1);
    }

    public synchronized void limitNextBytes(int len) {
        //累计已经发送/接收了多少字节数
        this.bytesWillBeSentOrReceive += len;
        //当累计的字节数大于定义的数据块大小时
        while (this.bytesWillBeSentOrReceive > CHUNK_LENGTH) {
            long nowTick = System.nanoTime();
            //计算积累数据期间消耗的时间
            long passTime = nowTick - this.lastPieceSentOrReceiveTick;
            //timeCostPerChunk表示单个块最多需要多少纳秒
            //如果missedTime大于0，说明此时流量进出的速率已经超过maxRate了，需要休眠来限制流量
            long missedTime = this.timeCostPerChunk - passTime;
            if (missedTime > 0) {
                try {
                    Thread.sleep(missedTime / 1000000, (int) (missedTime % 1000000));
                } catch (InterruptedException e) {
                    LOGGER.error(e.getMessage(), e);
                }
            }
            //重置累计发送/接收的字节数
            this.bytesWillBeSentOrReceive -= CHUNK_LENGTH;
            //重置最后一次检查时间
            this.lastPieceSentOrReceiveTick = nowTick + (missedTime > 0 ? missedTime : 0);
        }
    }
}

有了限流器后，现在我们要对下载功能做限流。因为 java 的 io 流的设计是装饰器模式，因此我们可以方便的封装一个我们自己的 InputStream

package com.example.demo.limit;
import java.io.IOException;
import java.io.InputStream;

public class LimitInputStream extends InputStream {
    private InputStream inputStream;
    private BandwidthLimiter bandwidthLimiter;

    public LimitInputStream(InputStream inputStream, BandwidthLimiter bandwidthLimiter) {
        this.inputStream = inputStream;
        this.bandwidthLimiter = bandwidthLimiter;
    }

    @Override
    public int read() throws IOException {
        if (bandwidthLimiter != null) {
            bandwidthLimiter.limitNextBytes();
        }
        return inputStream.read();
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
        if (bandwidthLimiter != null) {
            bandwidthLimiter.limitNextBytes(len);
        }
        return inputStream.read(b, off, len);
    }

    @Override
    public int read(byte[] b) throws IOException {
        if (bandwidthLimiter != null && b.length > 0) {
            bandwidthLimiter.limitNextBytes(b.length);
        }
        return inputStream.read(b);
    }
}

后面我们使用这个 LimitInputStream 来读取文件，每次读取一块数据，限流器都会检查当前的速率是否超过指定的最大速率。这样就能间接的达到限制下载速率的目的了。

附上 SpringMVC 的一个下载限流的 demo

package com.example.demo.limit;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
@Slf4j
@RestController
public class limitDownloadFile {

    @GetMapping("/limit")
    public void limitDownloadFile(String file, HttpServletResponse response) throws IOException {
        log.info("download file");
        if (file == null) {
            file = "D:\\KK录像机\\KKcapture\\KKCamera.exe";
        }
        File downloadFile = new File(file);
        FileInputStream fileInputStream = new FileInputStream(downloadFile);

        response.setContentType("application/x-msdownload;");
        response.setHeader("Content-disposition", "attachment; filename=" + new String(downloadFile.getName()
                .getBytes("utf-8"), "ISO8859-1"));
        response.setHeader("Content-Length", String.valueOf(downloadFile.length()));
        ServletOutputStream outputStream = null;
        try {
            LimitInputStream limitInputStream = new LimitInputStream(fileInputStream, new BandwidthLimiter(100));

            long beginTime = System.currentTimeMillis();
            outputStream = response.getOutputStream();
            byte[] bytes = new byte[1024];
            int read = limitInputStream.read(bytes, 0, 1024);
            while (read != -1) {
                outputStream.write(bytes);
                read = limitInputStream.read(bytes, 0, 1024);
            }
            log.info("download use {} ms", System.currentTimeMillis() - beginTime);
        } finally {
            fileInputStream.close();
            if (outputStream != null) {
                outputStream.close();
            }
            log.info("download success!");
        }
    }
}

三、注意点

使用这个算法要注意一个问题，就是 chunk 的块大小不能设置的太小，即 CHUNK_LENGTH 不能设置的太小。否则容易造成明明 maxRate 设置的很大，但是实际下载速率却很小的问题。

假设 CHUNK_LENGTH 就设置为 1024 bytes，每次读取的块大小也是 1024 bytes，maxRate 为 64M/s。那么我们可以计算出 timeCostPerChunk 约等于 15258 纳秒。

再如果真正的速率是 100M/s，也就是每秒差不多会调用 limitNextBytes 方法 100000 次，由于每次读取消耗的时间极短，因此每次进入该方法都要 sleep 15258 纳秒之后再读取下一个块的数据。** 如果没有算上线程调度的时间，就算 1 秒内休眠 100000 次也完全没什么问题。** 但是线程的休眠和唤醒都需要内核来进行，线程上下文切换的时间应该远大于 15258 纳秒，这时候频繁的休眠就会导致线程暂停运行的时间和我们预期的不符。由于休眠时间过长，最终导致实际的下载速率大大的低于 maxRate。

因此，我们需要调大 CHUNK_LENGTH，尽量让 timeCostPerChunk 的值远大于线程调度的时间，减少线程调度对限流造成的影响。

文章转载：https://my.oschina.net/lenglingx/blog/5573262