简单实现令牌桶限流

asmall_cat2023-09-12 21:03

java 复制代码

package com.trpc.protection;

import lombok.extern.slf4j.Slf4j;

/**
 * 令牌桶限流器
 */
@Slf4j
public class TokenBucketRateLimiter implements RateLimiter{
    //令牌桶 >0有令牌
    private int tokenBucket;

    //令牌桶最大容量
    private final int max;

    //加入令牌的速率
    private final int rate;

    //上次放入令牌时间
    private long lastPutTokenTime;

    public TokenBucketRateLimiter( int max, int rate) {
        this.tokenBucket = max;
        this.max = max;
        this.rate = rate;
        this.lastPutTokenTime = System.currentTimeMillis();
    }
    /**
     * 判断是否放行请求
     * 多服务调用限流-->加锁
     * @return
     */
    public synchronized boolean allowRequest(){

        //1 添加令牌
        long currentTime = System.currentTimeMillis();  //拿到当前时间

        long needAddTokenTime = currentTime - lastPutTokenTime; //拿到时间差

        //2 设置间隔一秒放一次令牌
        if (needAddTokenTime >= 1000 ){
            int addTokenNums= (int) (needAddTokenTime * rate / 1000);// 乘上速率拿到需要添加的令牌书

            tokenBucket = Math.min(max, tokenBucket + addTokenNums);//给令牌桶添加令牌，到达最大值就用最大值，否则就用桶内余量+新增余量

            this.lastPutTokenTime=System.currentTimeMillis();//重置
        }

        //2自增令牌
        if (tokenBucket > 0){
            tokenBucket--;
            if (log.isDebugEnabled()){
                log.debug("------还有【{}】令牌存在，放行--------",tokenBucket);
            }
            return true;
        }else {
            if (log.isDebugEnabled()){
                log.debug("限流中请等待");
            }
            return false;
        }
    }
}

测试

java 复制代码

public static void main(String[] args) {
        TokenBucketRateLimiter tokenBucketRateLimiter = new TokenBucketRateLimiter(20,20);
        for (int i = 0; i < 21; i++) {
            try {
                Thread.sleep(10);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            boolean b = tokenBucketRateLimiter.allowRequest();
            System.out.println("是否放行 = " + b);
        }
    }