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

r囧r小猫 2023-03-01 06:15 14阅读 0赞

首先准备一些必要的参数

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

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

[完整源码  
][Link 1]

[Link 1]: https://github.com/AmyliaScarlet/TokenBucket