com.ghp.request.limiter.impl.SlideWindowCountLimiter Maven / Gradle / Ivy

Go to download
package com.ghp.request.limiter.impl;

import com.ghp.request.annotation.RequestLimit;
import com.ghp.request.limiter.Limiter;

import java.util.LinkedList;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLongArray;

/**
 * @author ghp
 * @title 滑动窗口计数限流器
 * @description
 */
public class SlideWindowCountLimiter implements Limiter {

    /**
     * 映射接口的计数器
     * key 为接口的名称，value是接口对应的计数器
     */
    private Map counterMap = new ConcurrentHashMap<>();

    private class SlidingWindowCounter {
        /**
         * 每秒可以处理的请求数量（也就是 QPS）
         */
        private final int qps;
        /**
         * 滑动窗口的大小
         */
        private final int interval;
        /**
         * 记录每一个滑块的请求数量
         */
        private LinkedList sliders;
        /**
         * 记录每一个滑块的时间戳
         */
        private LinkedList timestamps;
        /**
         * 记录窗口中请求的总数量
         */
        private int count;
        /**
         * 滑块的数量
         */
        private final int sliderNumber;


        public SlidingWindowCounter(int qps, int interval, int slider) {
            this.qps = qps;
            this.interval = interval;
            this.sliderNumber = interval / slider;
            this.sliders = new LinkedList<>();
            this.timestamps = new LinkedList<>();
        }

        /**
         * 判断是否允许请求
         *
         * @return true-允许 false-不允许
         */
        public boolean allowRequest() {
            long now = System.currentTimeMillis();

            // 移除过期的滑块
            for (int i = 0; i < sliders.size(); i++) {
                long timestamp = timestamps.isEmpty() ? 0 : timestamps.get(i);
                if (now - timestamp <= interval) {
                    // 由于sliders中是按照时间戳排序的，当发现第一个滑块未过期，后面的就都不会过期
                    break;
                }
                count -= sliders.removeFirst();
                timestamps.removeFirst();
            }

            // 判断请求数量是否超过阈值
            if ((count + 1) > qps) {
                return false;
            }
            // 未超过阈值，更新窗口中的滑块，已经当前窗口的最大请求数
            int lastCount = sliders.isEmpty() ? 0 : sliders.peekLast();
            if (sliders.size() < sliderNumber) {
                // 窗口中滑块数量未满，则直接新增滑块
                sliders.addLast(1);
                timestamps.addLast(now);
            } else {
                // 窗口中滑块数量已满，则直接更新最后一个滑块
                sliders.set(sliderNumber - 1, lastCount + 1);
                timestamps.set(sliderNumber - 1, now);
            }
            count++;
            return true;
        }
    }


    /**
     * 限流
     *
     * @param methodName
     * @param requestLimitAnnotation
     * @return true-发生限流 false-未限流
     */
    @Override
    public boolean limit(String methodName, RequestLimit requestLimitAnnotation) {
        int qps = requestLimitAnnotation.qps();
        int interval = requestLimitAnnotation.interval();
        int slider = requestLimitAnnotation.slider();
        SlidingWindowCounter slidingWindowCounter = counterMap.computeIfAbsent(methodName, key ->
                new SlidingWindowCounter(qps, interval, slider));
        return !slidingWindowCounter.allowRequest();
    }
}