com.ghp.request.limiter.impl.TokenBucketLimiter 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 TokenBucketLimiter implements Limiter {

    /**
     * 记录上一次请求时漏桶的容量
     * 如果初始容量为 10，表示一个桶最多同时处理10个请求，超出部分就会被限流
     * key是接口名，value是漏桶的剩余容量
     */
    private Map tokenCountMap = 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 preTokenCount = tokenCountMap.getOrDefault(methodName, capacity);
        // 上一次请求时的时间戳（初始状态为当前时间）
        long prevTimestamp = this.timestampMap.getOrDefault(methodName, now);
        // 上一次请求与当前请求的时间间隔
        int interval = (now - prevTimestamp) == 0 ? 0 : (int) ((now - prevTimestamp) / 1000);

        int curTokenCount = preTokenCount + interval * qps;
        curTokenCount = Math.min(curTokenCount, capacity);
        // 判断当前桶的容量是否足够
        if (curTokenCount < 1) {
            // 当前桶中的令牌数量不足1，需要限流
            return true;
        }
        // 当前容量令牌数量足够，则更新令牌数和时间戳
        timestampMap.put(methodName, now);
        tokenCountMap.put(methodName, curTokenCount - 1);
        return false;
    }
}