滑动窗口版限流算法-----------线程安全版写法

上一篇文章贴了滑动窗口版限流算法单线程版，也就是再单线程的情况下是正常运行的，那么这篇文章就贴一下在多线程的情况下线程安全的写法，主要运用的是 cas，以及JUC的一些原子类，代码内部有注释，这边就不详细说明了，如果有错误的地方，欢迎指正!

package com.ccc;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: chencc
 */
public class ConcurrentSlidingWindowRateLimit {
    /**
     * 窗口大小
     */
    private int windowSize;
    /**
     * 最大请求数
     */
    private int maxRequestNum;
    /**
     * 以1000毫秒为一个分割(qps)
     */
    private final int LIMIT_UNIT = 1000;
    /**
     * 窗口数组 模拟队列
     */
    private Window[] windowArr;
    /**
     * 窗口数组的长度
     */
    private int arrLen;
    /**
     * 尾下标
     */
    private AtomicInteger tailIndex;
    /**
     * 头下标
     */
    private AtomicInteger headIndex;
    /**
     * 当前请求数
     */
    private AtomicInteger currentRequestNum;

    class Window {
        /**
         * 窗口的起点
         */
        private long start;
        /**
         * 在这个区间内的请求数
         */
        private AtomicInteger requestNum;
        /**
         * 当前窗口的状态 1激活状态  -1无效状态
         */
        private AtomicInteger state;

        public Window() {
            this.requestNum = new AtomicInteger(0);
            this.state = new AtomicInteger(-1);
        }

        public Window(int start, int requestNum) {
            this.start = start;
            this.requestNum = new AtomicInteger(0);
        }
    }

    public ConcurrentSlidingWindowRateLimit(int maxRequestNum) {
        this(100, maxRequestNum);
    }

    public ConcurrentSlidingWindowRateLimit(int windowSize, int maxRequestNum) {
        if (LIMIT_UNIT % windowSize != 0) {
            throw new IllegalArgumentException(LIMIT_UNIT + " % windowSize must is 0");
        }
        this.windowSize = windowSize;
        this.maxRequestNum = maxRequestNum;
        this.arrLen = getLen(LIMIT_UNIT / windowSize);
        this.windowArr = new Window[this.arrLen];
        tailIndex = new AtomicInteger(0);
        headIndex = new AtomicInteger(0);
        currentRequestNum = new AtomicInteger(0);
        // 预先初始化window对象
        initWindowArr(this.arrLen);
        // 方便位运算取模的时候不用再减一了
        this.arrLen--;
    }

    private void initWindowArr(int len) {
        for (int n = 0; n < len; n++) {
            windowArr[n] = new Window();
        }
    }

    /**
     * 判断当前这个请求是否能够通过
     * @return
     */
    public boolean entry() {
        // 获得当前时间
        long currentTime = System.currentTimeMillis();
        int index;
        // 判断窗口是否可以移出去
        while (enableRemoveWindow(currentTime, index = tailIndex.get())) {
            // 删除窗口
            removeTailWindow(index);
        }
        return doEntry(currentTime);
    }

    private boolean doEntry(long time) {
        int index, count = 0;
        for (;currentRequestNum.get() < maxRequestNum;) {
            // 如果头下标等于尾下标
            if ((index = headIndex.get()) == tailIndex.get()) {
                activateWindow(time, index);
            }
            // 获得当前最新的window对象下标索引
            int headWindowIndex = getCurrentHeadWindowIndex();
            Window currentHeadWindow = windowArr[headWindowIndex];
            // 如果处于未激活状态 自旋等待
            if (currentHeadWindow.state.get() < 0) {
                continue;
            }
            // 如果window的开始时间已经大于等于当前这个时间
            // 如果count == 0 尝试更新一次时间
            // 如果count == 1 代表已经更新过一次了 可能在里面 自旋太久了 直接返回false
            if (currentHeadWindow.start > time) {
                if (count == 0) {
                    time = System.currentTimeMillis();
                    count++;
                    continue;
                }
                return false;
            }
            // 如果不需要新生成一个窗口
            if (!check(currentHeadWindow, time)) {
                int requestNum;
                // 如果当前请求数量小于最大请求量  并且 窗口状态正常 一直尝试cas 直到成功或者条件不成立
                for (;(requestNum = currentRequestNum.get()) < maxRequestNum && currentHeadWindow.start <= time && currentHeadWindow.state.get() > 0;) {
                    // 如果cas成功的话
                    if (currentRequestNum.compareAndSet(requestNum, requestNum + 1)) {
                        currentHeadWindow.requestNum.incrementAndGet();
                        return true;
                    }
                }
            } else {
                // 如果当前时间已经和最新的那个窗口相差了至少一个窗口大小 需要新激活一个窗口 再次检查，防止中间已经变化了
                if (headWindowIndex == getCurrentHeadWindowIndex() && check(currentHeadWindow = windowArr[headWindowIndex], time)) {
                    activateWindow(getStartTime(currentHeadWindow, time), (headWindowIndex + 1) & arrLen);
                }
            }
        }
        return false;
    }

    /**
     * 当这个窗口和这个时间不在一个窗口内，获得这个时间应该在的窗口的startTime(可能中间时间已经相差了好几个窗口)
     * @param window
     * @param time
     * @return
     */
    public long getStartTime(Window window, long time) {
        long start = window.start;
        long interval = time - start;
        // 获得相差窗口的个数
        long num = interval / windowSize;
        return num * windowSize + start;
    }

    /**
     * 检查是否需要新生成一个窗口，也就是这个时间是否在这个窗口内
     * @param currentHeadWindow
     * @param time
     * @return
     */
    private boolean check(Window currentHeadWindow, long time) {
        return currentHeadWindow.start + windowSize <= time;
    }

    /**
     * 获得当前最新的window对象的下标索引
     * @return
     */
    private int getCurrentHeadWindowIndex() {
        if (headIndex.get() == 0) {
            return arrLen;
        }
        return headIndex.get() - 1;
    }

    /**
     * 激活对应下标窗口
     * @param startTime
     * @param index
     */
    private void activateWindow(long startTime, int index) {
        Window window = windowArr[index];
        // 可能这个窗口的状态不是为失效的状态
        if (window.state.get() == 1) {
            return;
        }
        if (headIndex.compareAndSet(index, (index + 1) & arrLen)) {
            // 激活窗口
            window.requestNum.set(0);
            window.start = startTime;
            window.state.set(1);
        }
        return;
    }

    /**
     * 移动尾下标 来表示移除掉窗口
     */
    private void removeTailWindow(int index) {
        // 如果cas成功的话
        if (tailIndex.compareAndSet(index, (index + 1) & arrLen)) {
            Window window = windowArr[index];
            currentRequestNum.addAndGet(-window.requestNum.get());
            window.state.set(-1);
        }
    }

    /**
     * 判断是否能够移除掉窗口
     * @param time
     * @return
     */
    private boolean enableRemoveWindow(long time, int index) {
        if (index == headIndex.get()) {
            return false;
        }
        return time - windowArr[index].start >= LIMIT_UNIT;
    }

    /**
     * 获得大于len的最小2的N次方的数
     * 参考 ArrayDeque 里面的 源码 基本拷贝的 哈哈
     * 为了 % 可以使用位运算
     * @param len
     * @return
     */
    private int getLen(int len) {
        len |= (len >>>  1);
        len |= (len >>>  2);
        len |= (len >>>  4);
        len |= (len >>>  8);
        len |= (len >>> 16);
        len++;
        if (len < 0) {
            len >>>= 1;
        }
        return len;
    }
}