Springboot实现令牌桶限流

在对一个秒杀系统进行接口访问限流，感觉计数法太过于简单粗暴，便实现一个简单的令牌桶限流

令牌桶

大小固定的令牌桶可自行以恒定的速率源源不断地产生令牌。如果令牌不被消耗，或者被消耗的速度小于产生的速度，令牌就会不断地增多，直到把桶填满。后面再产生的令牌就会从桶中溢出。最后桶中可以保存的最大令牌数永远不会超过桶的大小。

---

设计思路

• 实现令牌桶类
• 用HashMap<>存储多个桶，为各个接口提供服务，静态便于访问
• 记得要重写hash()和equal()方法
• 配置Spring Boot计时器，实现按一定速率生产令牌
• 每次访问接口时，取走一个令牌

---

令牌桶实现类

import java.util.HashMap;
import java.util.Objects;

public class Bucket {

    //默认容量10 
    static final int DEFAULT_MAX_COUNT = 10;
    // 默认增长速率为1
    static final int DEFAULT_CREATE_RATE = 1;
    public static HashMap<String, Bucket> buckets = new HashMap(10);

    final String name;
    final int maxCount;
    //当前令牌数
    int size;
    int createRate;

    public Bucket() {
        maxCount = DEFAULT_MAX_COUNT;
        createRate = DEFAULT_CREATE_RATE;
        name = buckets.size() + "bucket";
    }

    public Bucket(String name, int maxCount, int createRate) {
        this.name = name;
        this.maxCount = maxCount;
        this.createRate = createRate;
    }

    @Override
    public boolean equals(Object obj) {
        if (obj == null)
            return false;
        Bucket bucket = (Bucket) obj;
        if (bucket.size != size || bucket.createRate != createRate || bucket.maxCount != maxCount)
            return false;
        return true;
    }

    @Override
    public int hashCode() {
        return Objects.hash(maxCount, size, createRate);
    }

    public String getName() {
        return name;
    }

    public int getSize() {
        return size;
    }

    public boolean isFull() {
        return size == maxCount;
    }

    //根据速率自增
    public void incrTokens() {
        for (int i = 0; i < createRate; i++)
            incrOneToken();
    }

    //生产一个令牌
    void incrOneToken() {
        if (isFull())
            return;
        size++;
    }

    //消费一个
    public boolean getToken() {
        if (size > 0)
            size--;
        else
            return false;
        return true;
    }

}

---

启动类里配置计时器

启动类加入注解@EnableScheduling

按一定速率生产令牌

 @Scheduled(fixedRate = 1000)
    public void timer() {
        //获取当前时间
        LocalDateTime localDateTime = LocalDateTime.now();
        System.out.println("当前时间为:" + localDateTime.format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        if (Bucket.buckets.containsKey("seckill"))
            Bucket.buckets.get("seckill").incrTokens();
    }

---

调用接口时取令牌

//如果取到执行
if(Bucket.buckets.get("seckill").getToken()){
                //do something
            }
else{
    //
}