【Java】Spring自定义注解令牌桶

首先准备一些必要的参数

public class Constants {
    //容量
    public static final int CAPACITY = 10;
    //增长速率
    public static final int RISE_COUNT = 1;
    //增长时间间隔(毫秒)
    public static final int fixedRate = 1000;
    //所有的桶名
    public static final String[] BUCKETS = {"demo","demo2","demo3"};
}

构造一个桶类，桶要能生成令牌，存放令牌，取出令牌

public class Bucket {

    private final String name;
    private final int maxCount;

    private int size;
    private final int riseCount;

    public Bucket() {
        maxCount = Constants.CAPACITY;
        riseCount = Constants.RISE_COUNT;
        name = "bucket";
    }
    public Bucket(String name) {
        this.maxCount = Constants.CAPACITY;
        this.riseCount = Constants.RISE_COUNT;
        this.name = name;
    }
    public Bucket(String name, int maxCount, int riseCount) {
        this.name = name;
        this.maxCount = maxCount;
        this.riseCount = riseCount;
    }

    @Override
    public boolean equals(Object obj) {
        if (obj == null)
            return false;
        Bucket bucket = (Bucket) obj;
        return bucket.size == size && bucket.riseCount == riseCount && bucket.maxCount == maxCount;
    }
   
    @Override
    public int hashCode() {
        return Objects.hash(name, maxCount, size, riseCount);
    } 
 
    public String getName() {
        return name;
    }

    public int getSize() {
        return size;
    }

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

    public void genTokens() {
        for (int i = 0; i < riseCount; i++)
            genToken();
    }

    void genToken() {
        if (!isFull())
            size++;
    }

    public synchronized int getToken() {
        if (size > -1)
            size--;
        return size;
    }

}

创建BucketManager

public class BucketManager {
    private static final BucketManager instance = new BucketManager();

    public static volatile HashMap<String, Bucket> buckets = new HashMap(Constants.CAPACITY);

    public static BucketManager getInstance(){
        return instance;
    }

    public void genBucket(String name){
        Bucket bucket = new Bucket(name);
        buckets.put(name,bucket);
    }

    public synchronized HashMap<String, Bucket> getBuckets(){
        return buckets;
    }
    public synchronized Bucket getBucket(String name) throws BucketManagerException {
        if(buckets.containsKey(name)){
            return buckets.get(name);
        }
        throw new BucketManagerException("not contain key: "+ name);
    }

启动类
添加@EnableScheduling

@EnableScheduling
@SpringBootApplication
public class TokenBucketApplication {

	public static void main(String[] args) {
		SpringApplication.run(TokenBucketApplication.class, args);

		//创建令牌桶
		for (String name: Constants.BUCKETS)
			BucketManager.getInstance().genBucket(name);

	}

	@Scheduled(fixedRate = Constants.fixedRate)
	public void timer() {
		BucketScheduled.genTokens(); //定时生产令牌
	}


}

这样令牌桶就完成了。

接下来创建自定义注解

/**
 * 需要消耗令牌的方法加上该注释 设置桶名
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RequireToken {
    /**桶名*/
    String BucketName();
}

创建拦截器TokenBucketInterceptor

public class TokenBucketInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws BucketManagerException {
        //不相干的
        if (!(handler instanceof HandlerMethod)) {
            return true;
        }

        boolean isAssignableFrom = handler.getClass().isAssignableFrom(HandlerMethod.class);
        //有注解的
        if(isAssignableFrom){
            RequireToken assertTokenBucket = ((HandlerMethod) handler).getMethodAnnotation(RequireToken.class);
            //有RequireToken注解的
            if(assertTokenBucket == null || assertTokenBucket.BucketName().equals("")){
                return true;
            }else {

                int token = BucketManager.getInstance().getBucket(assertTokenBucket.BucketName()).getToken();
                if(token>=0){
                    System.out.println("Bucket:"+ assertTokenBucket.BucketName() +" get token success,surplus:"+token);
                    return true;
                }
                else{
                    System.out.println("Bucket:"+ assertTokenBucket.BucketName()+" none token");
                    return false;
                }
            }
        }
        return true;
    }
}

WebMvcConfigurer添加拦截器

@Configuration
public class WebMvcConfigurer extends WebMvcConfigurationSupport {

   @Override
   public void addInterceptors(InterceptorRegistry registry)
    {
        registry.addInterceptor(tokenBucketInterceptor()).addPathPatterns("/**");
        super.addInterceptors(registry);
    }

    @Bean
    public TokenBucketInterceptor tokenBucketInterceptor()
    {
        return new TokenBucketInterceptor();
    }
}

注解创建完成了。

在需要使用令牌桶的接口方法上加上注解@RequireToken

@RestController
@ResponseBody
public class Api {

    @RequireToken(BucketName="demo")
    @GetMapping("/demo")
    public String demo() {
        System.out.println("demo dosomething");
        return "ok";
    }
}

以1秒生成一个令牌，容量为10 举例
快速请求11次
输出：
Bucket:demo get token success,surplus:9
demo dosomething
Bucket:demo get token success,surplus:8
demo dosomething
Bucket:demo get token success,surplus:7
demo dosomething
Bucket:demo get token success,surplus:6
demo dosomething
Bucket:demo get token success,surplus:5
demo dosomething
Bucket:demo get token success,surplus:4
demo dosomething
Bucket:demo get token success,surplus:3
demo dosomething
Bucket:demo get token success,surplus:2
demo dosomething
Bucket:demo get token success,surplus:1
demo dosomething
Bucket:demo get token success,surplus:0
demo dosomething
Bucket:demo none token

以上是一个基础版
根据业务需要，令牌的生产及管理可以与桶分开，且根据应对高峰情况对容量和增量做动态调整。

完整源码