Java限流代码实现

常见的限流算法大致有三种：

• 令牌桶算法
• 漏桶算法
• 计数器算法

网上对令牌桶又细分为固定窗口计数器限流和滑动窗口计数器限流，下面将对这几种限流方式进行简单的介绍及代码实现。
注意：代码中会考虑并发线程安全问题，非分布式限流
Github地址：重构后的代码

固定窗口计数器限流

固定窗口计数器限流就是在固定时间内（如10s），只允许固定的请求数访问（如10个），超过的请求将受到限制。
实现逻辑图

实现代码

package com.dfy.ratelimiter.core;

import java.util.concurrent.TimeUnit;

/**
 * @description: 计数器限流
 * @author: DFY
 * @time: 2020/4/8 17:02
 */
public abstract class CounterLimit {

    /** 单位时间限制数 */
    protected int limitCount;
    /** 限制时间 */
    protected long limitTime;
    /** 时间单位，默认为秒 */
    protected TimeUnit timeUnit;

    /** 当前是否为受限状态 */
    protected volatile boolean limited;

    /**
     * 尝试将计数器加1，返回为true表示能够正常访问接口，false表示访问受限
     * @return
     */
    protected abstract boolean tryCount();
}

package com.dfy.ratelimiter.core;

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

import java.time.LocalDateTime;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @description: 固定窗口计数器限流
 * @author: DFY
 * @time: 2020/4/8 15:50
 */
public class FixedWindowCounterLimit extends CounterLimit {

    private static Logger logger = LoggerFactory.getLogger(FixedWindowCounterLimit.class);

    /** 计数器 */
    private AtomicInteger counter = new AtomicInteger(0);

    public FixedWindowCounterLimit(int limitCount, long limitTime) {
        this(limitCount, limitTime, TimeUnit.SECONDS);
    }

    public FixedWindowCounterLimit(int limitCount, long limitTime, TimeUnit timeUnit) {
        this.limitCount = limitCount;
        this.limitTime = limitTime;
        this.timeUnit = timeUnit;
        new Thread(new CounterResetThread()).start(); // 开启计数器清零线程
    }

    public boolean tryCount() {
        while (true) {
            if (limited) {
                return false;
            } else {
                int currentCount = counter.get();
                if (currentCount == limitCount) {
                    logger.info("限流：{}", LocalDateTime.now().toString());
                    limited = true;
                    return false;
                } else {
                    if (counter.compareAndSet(currentCount, currentCount + 1))
                        return true;
                }
            }
        }
    }

    class CounterResetThread implements Runnable {
        @Override
        public void run() {
            while (true) {
                try {
                    timeUnit.sleep(limitTime);
                    counter.compareAndSet(limitCount, 0); // 计数器清零
                    limited = false; // 修改当前状态为不受限
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}

使用及测试
启动项目，连续访问接口，当在访问第11次时接口受限，受限时间到后又能正常访问。

private FixedWindowCounterLimit fixedWindowCounterLimit = new FixedWindowCounterLimit(10, 10);
@GetMapping("/hello")
public String hello() {
    if (!fixedWindowCounterLimit.tryCount()) {
        return "限流!";
    }
    return "hello world!";
}

存在的问题
限流不均匀，如下所示我们规定10S内至多10个访问量，但2S内实际上有20个访问量。

滑动窗口计数器限流

固定窗口计数器限流是在固定时间内访问量受限，滑动窗口计数器限流是在滑动窗口内访问量受限。
例子
如下是规定5S内不能超过10个访问量，当已经达到10个访问量，则访问受限。使用该方式可以使受限均匀，任意连续的5S内都只能有10个访问量。

实现代码

package com.dfy.ratelimiter.core;

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

import java.time.LocalDateTime;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @description: 滑动窗口计数器限流
 * @author: DFY
 * @time: 2020/4/8 17:01
 */
public class SlidingWindowCounterLimit extends CounterLimit {

    private static Logger logger = LoggerFactory.getLogger(SlidingWindowCounterLimit.class);

    /** 格子分布 */
    private AtomicInteger[] gridDistribution;
    /** 当前时间在计数分布的索引 */
    private volatile int currentIndex;
    /** 当前时间之前的滑动窗口计数 */
    private int preTotalCount;
    /** 格子数 */
    private int gridNumber;
    /** 是否正在执行状态重置 */
    private volatile boolean resetting;

    public SlidingWindowCounterLimit(int gridNumber, int limitCount, long limitTime) {
        this(gridNumber, limitCount, limitTime, TimeUnit.SECONDS);
    }

    public SlidingWindowCounterLimit(int gridNumber, int limitCount, long limitTime, TimeUnit timeUnit) {
        if (gridNumber <= limitTime)
            throw new RuntimeException("无法完成限流，gridNumber必须大于limitTime，gridNumber = " + gridNumber + ",limitTime = " + limitTime);
        this.gridNumber = gridNumber;
        this.limitCount = limitCount;
        this.limitTime = limitTime;
        this.timeUnit = timeUnit;
        gridDistribution = new AtomicInteger[gridNumber];
        for (int i = 0; i < gridNumber; i++) {
            gridDistribution[i] = new AtomicInteger(0);
        }
        new Thread(new CounterResetThread()).start();
    }

    public boolean tryCount() {
        while (true) {
            if (limited) {
                return false;
            } else {
                int currentGridCount = gridDistribution[currentIndex].get();
                if (preTotalCount + currentGridCount == limitCount) {
                    logger.info("限流：{}", LocalDateTime.now().toString());
                    limited = true;
                    return false;
                }
                if (!resetting && gridDistribution[currentIndex].compareAndSet(currentGridCount, currentGridCount + 1))
                    return true;
            }
        }
    }

    class CounterResetThread implements Runnable {
        @Override
        public void run() {
            while (true) {
                try {
                    timeUnit.sleep(1); // 停止1个时间单位
                    int indexToReset = currentIndex - limitCount - 1; // 要重置计数的格子索引
                    if (indexToReset < 0) indexToReset += gridNumber;
                    resetting = true; // 防止在更新状态时，用户访问接口将当前格子的访问量 + 1
                    preTotalCount = preTotalCount - gridDistribution[indexToReset].get()
                            + gridDistribution[currentIndex++].get(); // 重置当前时间之前的滑动窗口计数
                    if (currentIndex == gridNumber) currentIndex = 0;
                    if (preTotalCount + gridDistribution[currentIndex].get() < limitCount)
                        limited = false; // 修改当前状态为不受限
                    resetting = false;
                    logger.info("当前格子：{}，重置格子：{}，重置格子访问量：{}，前窗口格子总数：{}",
                            currentIndex, indexToReset, gridDistribution[indexToReset].get(), preTotalCount);
                    gridDistribution[indexToReset].set(0);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}

使用及测试

private SlidingWindowCounterLimit slidingWindowCounterLimit = new SlidingWindowCounterLimit(20, 10, 10);
@GetMapping("/hello")
public String hello() {
    if (!slidingWindowCounterLimit.tryCount()) {
        return "限流!";
    }
    return "hello world!";
}

令牌桶限流

Google guava的RateLimiter提供了基于令牌桶算法的两种实现，下面代码只是简单实现。

package com.dfy.ratelimiter.core;

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

import java.time.LocalDateTime;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @description: 令牌桶限流
 * @author: DFY
 * @time: 2020/4/10 15:35
 */
public class TokenBucketLimit {

    private static Logger logger = LoggerFactory.getLogger(TokenBucketLimit.class);

    /** 给定时间生成令牌数 */
    private int genNumber;
    /** 生成令牌所花费的时间 */
    private int genTime;
    /** 时间单位，默认为秒 */
    private TimeUnit timeUnit;
    /** 最大令牌数 */
    private int maxNumber;
    /** 已存储的令牌数 */
    private AtomicInteger storedNumber;

    public TokenBucketLimit(int genNumber, int genTime, int maxNumber) {
        this(genNumber, genTime, TimeUnit.SECONDS, maxNumber);
    }

    public TokenBucketLimit(int genNumber, int genTime, TimeUnit timeUnit, int maxNumber) {
        this.genNumber = genNumber;
        this.genTime = genTime;
        this.timeUnit = timeUnit;
        this.maxNumber = maxNumber;
        this.storedNumber = new AtomicInteger(0);
        new Thread(new TokenGenerateThread()).start();
    }

    public boolean tryAcquire() {
        while (true) {
            int currentStoredNumber = storedNumber.get();
            if (currentStoredNumber == 0) {
                logger.info("限流：{}", LocalDateTime.now().toString());
                return false;
            }
            if (storedNumber.compareAndSet(currentStoredNumber, currentStoredNumber - 1)) {
                return true;
            }
        }
    }

    class TokenGenerateThread implements Runnable {
        @Override
        public void run() {
            while (true) {
                if (storedNumber.get() == maxNumber) {
                    logger.info("当前令牌数已满");
                    try { timeUnit.sleep(genTime); }
                    catch (InterruptedException e) { e.printStackTrace(); }
                } else {
                    int old =  storedNumber.get();
                    int newValue = old + genNumber;
                    if (newValue > maxNumber)
                        newValue = maxNumber;
                    storedNumber.compareAndSet(old, newValue);
                    logger.info("生成令牌数：{}，当前令牌数：{}", genNumber, newValue);
                    try { timeUnit.sleep(genTime); }
                    catch (InterruptedException e) { e.printStackTrace(); }
                }
            }
        }
    }
}

漏桶算法

漏桶限流的实现与令牌桶限流类似，只是一个是按固定速率增加，一个按固定速率减少。

package com.dfy.ratelimiter.core;

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

import java.time.LocalDateTime;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @description: 漏桶限流
 * @author: DFY
 * @time: 2020/4/13 14:47
 */
public class LeakyBucketLimit {

    private static Logger logger = LoggerFactory.getLogger(LeakyBucketLimit.class);

    /** 桶最大容量 */
    private int maxNumber;
    /** 时间单位，默认为秒 */
    private TimeUnit timeUnit;
    /** 泄露的数量 */
    private int leakNumber;
    /** 泄露的时间 */
    private int leakTime;
    /** 桶中剩余数量 */
    private AtomicInteger remainingNumber;

    public LeakyBucketLimit(int leakNumber, int leakTime, int maxNumber) {
        this(leakNumber, leakTime, TimeUnit.SECONDS, maxNumber);
    }

    public LeakyBucketLimit(int leakNumber, int leakTime, TimeUnit timeUnit, int maxNumber) {
        this.leakNumber = leakNumber;
        this.leakTime = leakTime;
        this.timeUnit = timeUnit;
        this.maxNumber = maxNumber;
        this.remainingNumber = new AtomicInteger(0);
        
    }

    public boolean tryAcquire() {
        while (true) {
            int currentStoredNumber = remainingNumber.get();
            if (currentStoredNumber == maxNumber) {
                logger.info("限流：{}", LocalDateTime.now().toString());
                return false;
            }
            if (remainingNumber.compareAndSet(currentStoredNumber, currentStoredNumber + 1)) {
                return true;
            }
        }
    }

    class LeakThread implements Runnable {
        @Override
        public void run() {
            while (true) {
                if (remainingNumber.get() == 0) {
                    logger.info("当前桶已空");
                    try { timeUnit.sleep(leakTime); }
                    catch (InterruptedException e) { e.printStackTrace(); }
                } else {
                    int old =  remainingNumber.get();
                    int newValue = old - leakNumber;
                    if (newValue < 0)
                        newValue = 0;
                    remainingNumber.compareAndSet(old, newValue);
                    logger.info("泄露：{}，当前：{}", leakNumber, newValue);
                    try { timeUnit.sleep(leakTime); }
                    catch (InterruptedException e) { e.printStackTrace(); }
                }
            }
        }
    }
}

如有问题，欢迎指正！