线程安全的计数器限流

网上的代码实现的计数器限流基本都不是线程安全的，如下限流器的实现：

设定时间内允许通过设置的请求数

public class CountLimitNotSafe {

    /**
     * 计数器归零
     */
    private final AtomicLong ZERO = new AtomicLong(0);
    /**
     * 计数器
     */
    private AtomicLong counter = ZERO;
    /**
     * 起始时间
     */
    private static long timestamp = System.currentTimeMillis();
    /**
     * 允许的请求书
     */
    private final long permitsPerSecond;

    public CountLimitNotSafe(long permitsPerSecond) {
        this.permitsPerSecond = permitsPerSecond;
    }

    public boolean tryAcquire() {
        long now = System.currentTimeMillis();
        // 如果当前时间已经过了1ms
        if (now - timestamp < 1L) {
            // 判断此时间间隔计数器是否大于设置的请求数
            if (counter.get() < permitsPerSecond) {
                counter.incrementAndGet();
                return true;
            } else {
                return false;
            }
        } else {
            // 重新设置起始时间，计数器归零
            // 在这打印一下进入这个else的时间，用于观测，如果打印时间相等，则证明多个线程进入了
            // 此处如果有多个线程进入，会重复设置时间起始时间和计数器，线程不安全
            System.out.println(now);
            counter = ZERO;
            timestamp = now;
            return false;
        }
    }

    public static void main(String[] args) {
        // 设置10个允许请求数
        CountLimitNotSafe countLimit = new CountLimitNotSafe(10);
        // 线程池
        ExecutorService executorService = Executors.newFixedThreadPool(10);
        // 模拟100个线程
        for (int i = 0; i < 100; i++) {
            executorService.submit(() -> {
                boolean result = countLimit.tryAcquire();
            });
        }
    }

运行结果如下：

重置计数起始时间：1612143072327
重置计数起始时间：1612143072328
重置计数起始时间：1612143072328
重置计数起始时间：1612143072328
重置计数起始时间：1612143072328
重置计数起始时间：1612143072329
重置计数起始时间：1612143072329
重置计数起始时间：1612143072330

因为线程不安全，多个线程进入else块，会重复设置时间和计数器，所以上面打印的时间有很多相同的。

优化后线程安全的实现为：

public class CountLimitSafe {

    private final AtomicLong ZERO = new AtomicLong(0);
    private AtomicLong counter = ZERO;
    private static long timestamp = System.currentTimeMillis();
    private final long permitsPerSecond;

    public CountLimitSafe(long permitsPerSecond) {
        this.permitsPerSecond = permitsPerSecond;
    }

    public boolean tryAcquire() {
        long now = System.currentTimeMillis();
        if (now - timestamp < 1) {
            return getAcquire();
        } else {
            // 加锁
            synchronized (CountLimitSafe.class){
                // 第一个拿到的锁会重置计数器
                // 后面进入else的线程，再次进行获取许可
                if (now - timestamp < 1) {
                    return getAcquire();
                } else {
                    System.out.println("重置计数起始时间：" + now);
                    counter = ZERO;
                    timestamp = now;
                    return false;
                }
            }
        }
    }

    private boolean getAcquire() {
        if (counter.get() < permitsPerSecond) {
            counter.incrementAndGet();
            return true;
        } else {
            return false;
        }
    }

    public static void main(String[] args) {
        CountLimitSafe countLimitSafe = new CountLimitSafe(10);
        ExecutorService executorService = Executors.newFixedThreadPool(30);
        for (int i = 0; i < 500; i++) {
            executorService.submit(() -> {
                boolean result = countLimitSafe.tryAcquire();
                if (result) {
                    System.out.println(result);
                }
            });
        }
    }
}

执行结果：

重置计数起始时间：1612143936234
重置计数起始时间：1612143936235
重置计数起始时间：1612143936236
重置计数起始时间：1612143936237
重置计数起始时间：1612143936238
重置计数起始时间：1612143936239
重置计数起始时间：1612143936240
重置计数起始时间：1612143936244
重置计数起始时间：1612143936245

无重复时间戳，说明没有重复重置计数器。

重置计数器的操作，放在在同步代码块，在进入同步代码块之后先判断，是否可以获取许可，由于第一个进入同步代码块的线程重置好了，后面进入的根据条件判断后，直接返回获取的结果，不会再重复设置计数器。跟单例模式的double check实现思路差不多一致。

这只是单机模式的计数器限流，在分布式微服务环境下，需要实现分布式限流。可以借助redis+lua脚本去实现，lua脚本中实现限流的整体逻辑，由于redis单线程的模型，lua脚本在执行的会顺序执行，所以不会发生线程不安全问题。