com.ghp.request.limiter.impl.LeakyBucketLimiter 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.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ghp
 * @title 漏桶限流器
 * @description
 */
public class LeakyBucketLimiter implements Limiter {

    /**
     * 记录上一次请求时漏桶的容量
     * 如果初始容量为 10，表示一个桶最多同时处理10个请求，超出部分就会被限流
     * key是接口名，value是漏桶的剩余容量
     */
    private Map capacityMap = new ConcurrentHashMap<>();
    /**
     * 记录接口的上一次请求起始时间
     * key是接口名，value是上一次请求的时间戳
     */
    private Map timestampMap = new ConcurrentHashMap<>();

    /**
     * 采用漏桶限流算法实现限流
     *
     * @param methodName
     * @param requestLimitAnnotation
     * @return true-发生限流 false-未限流
     */
    @Override
    public boolean limit(String methodName, RequestLimit requestLimitAnnotation) {
        int qps = requestLimitAnnotation.qps();
        int capacity = requestLimitAnnotation.capacity();

        long now = System.currentTimeMillis();
        // 上一次请求时桶的容量（初始状态为桶为空）
        int preCapacity = capacityMap.getOrDefault(methodName, 0);
        // 上一次请求时的时间戳（初始状态为当前时间）
        long prevTimestamp = this.timestampMap.getOrDefault(methodName, now);
        // 上一次请求与当前请求的时间间隔
        int interval = (now - prevTimestamp) == 0 ? 0 : (int) ((now - prevTimestamp) / 1000);

        // 当前桶的容量 = 上一次请求时桶的容量 - 时间间隔中漏出的水
        int curCapacity = preCapacity - interval * qps;
        curCapacity = Math.max(curCapacity, 0);
        // 判断当前桶的容量是否足够
        if ((curCapacity + 1) > capacity) {
            // 当前容量不够，需要限流
            return true;
        }
        // 当前容量足够，则更新容量和时间戳
        timestampMap.put(methodName, now);
        capacityMap.put(methodName, curCapacity + 1);
        return false;
    }
}