滑动窗口

138 阅读1分钟

public class SlidingWindowRateLimiter {

// 小窗口链表
LinkedList<Window> linkedList = null;
// 每个小窗口的时间间距
long interval4EachWindow = 1000;
// 小窗口数量
long windowCount = 10;
// 每个小窗口的限制
long countLimit4EachWindow = 110;
// 整个窗口的限制
long countLimit = 150;
// 实际请求数量
long count = 0;


/**
 *
 * @param countLimit 请求限制总数 如150
 * @param interval 时间段 如10s
 */
public SlidingWindowRateLimiter(long countLimit, long interval){
    this.interval4EachWindow = interval * 1000/ windowCount;
    this.countLimit = countLimit;
    // 每个小窗口的请求量限制数 设置为平均的2倍
    //this.countLimit4EachWindow = countLimit / windowCount * 2;
    // 时间窗口开始时间
    long start = System.currentTimeMillis();
    // 初始化连续的小窗口链表
    initWindow(start);
}


/**
 * 初始化窗口
 * @param start
 */
private void initWindow(long start){
    linkedList = new LinkedList<>();
    for (long i = 0; i < windowCount; i++) {
        linkedList.add(new Window(start, start += interval4EachWindow));
        //linkedList.add(new Window(start, start = start + i * interval4EachWindow));
    }
    // 总记数清零
    count = 0;
}


public boolean limit(){
    Window window = getAndRefreshWindows(System.currentTimeMillis());
    long windowRequestCount = window.getRequestCount();
    if(windowRequestCount + 1 <= countLimit4EachWindow && count + 1 <= countLimit){
        window.increase();
        count++;
        return true;
    }
    return false;
}


private Window getAndRefreshWindows(long requestTime){
    Window firstWindow = linkedList.getFirst();
    Window lastWindow = linkedList.getLast();
    // 发起请求时间在主窗口内
    if(requestTime > firstWindow.getStartTime()  && requestTime < lastWindow.getEndTime()){
        long distanceFromFirst = requestTime - firstWindow.getStartTime();
        long index = distanceFromFirst / interval4EachWindow;
        return linkedList.get((int) index);
    }else{
        long distanceFromLast = requestTime - lastWindow.getEndTime();
        long num = (long)(distanceFromLast / interval4EachWindow);
        // 请求时间超出主窗口一个窗口以上的身位
        if(num >= windowCount){
            initWindow(requestTime);
            return linkedList.getFirst();
        }else{
            slide(num+1);
            return linkedList.getLast();
        }
    }
}


/**
 * 移动窗口
 * @param stepNum
 */
private void slide(long stepNum){
    for (long i = 0; i < stepNum; i++) {
        Window removedWindow = linkedList.removeFirst();
        count = count - removedWindow.requestCount.get();
    }
    Window lastWindow = linkedList.getLast();
    long start = lastWindow.endTime;
    for (long i = 0; i < stepNum; i++) {
        linkedList.add(new Window(start, start += interval4EachWindow));
    }
}


public static void main(String[] args) throws InterruptedException {
    SlidingWindowRateLimiter slidingWindowRateLimiter = new SlidingWindowRateLimiter(300, 10);

    ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
    //每个500毫秒刷新一次窗口
    scheduledExecutorService.schedule(() -> {
        slidingWindowRateLimiter.getAndRefreshWindows(System.currentTimeMillis());
    }, 500, TimeUnit.MILLISECONDS);

    for (long i = 0; i < 400; i++) {
        System.out.println(slidingWindowRateLimiter.limit());
        Thread.sleep(5);
    }
}

class Window{
    
    //小窗口的开始时间
    @Getter
    private long startTime;
    //小窗口的结束时间
    @Getter
    private long endTime;
    //小窗口内的请求数量
    private AtomicLong requestCount;
    
    Window(long startTime, long endTime) {
        this.startTime = startTime;
        this.endTime = endTime;
        this.requestCount = new AtomicLong();
    }
    
    public long getRequestCount() {
        return requestCount.get();
    }
    
    public long increase(){
        return this.requestCount.incrementAndGet();
    }
}

}