Token bucket
In the case of high concurrency, current limiting is one of the common means at the back end. It can limit system current, interface current and user current. This paper uses token bucket algorithm + interceptor + custom annotation + custom exception to realize the demo of current limiting.
Token bucket idea
The token bucket with fixed size can continuously generate tokens at a constant rate. If the token is not consumed, or the consumed speed is less than the generated speed, the token will continue to increase until the bucket is filled.
The token generated later will overflow from the bucket. Finally, the maximum number of tokens that can be saved in the bucket will never exceed the size of the bucket. Then, each user who accesses will take a token from it, and can access only after getting the token. If the token is not obtained, it means that the access limit has been reached, and the access will be limited and not allowed
Implementation idea of current limiting demo
- Create token bucket class
- The project starts to initialize the token bucket, and sets a timer to regularly put tokens into the bucket
- Customize the current limiting annotation and mark the annotation on the interface that needs current limiting
- Configure the token bucket interceptor to intercept all paths, directly release the interfaces with infinite flow annotations, and handle the interfaces with limited flow annotations to get tokens. When a token is obtained, it will be released. If a token is not obtained, a user-defined exception will be thrown
- Customize exceptions and use AOP for global exception handling
In order to prevent concurrency, synchronized is added to the methods of generating and retrieving tokens
BucketUtil is as follows
1 public class BucketUtil { 2 3 //Default capacity 10 4 static final int DEFAULT_MAX_COUNT = 10; 5 // The default growth rate is 1 6 static final int DEFAULT_CREATE_RATE = 1; 7 // use HashMap Store token buckets. There are 10 token buckets by default 8 public static HashMap<String, BucketUtil> buckets = new HashMap(10); 9 10 //Custom capacity, once created, cannot be changed 11 final int maxCount; 12 //Custom growth rate 1 s Several tokens 13 int createRate; 14 //Current number of tokens 15 int size=0; 16 17 18 19 // Capacity and growth rate of default token bucket 20 public BucketUtil() { 21 maxCount = DEFAULT_MAX_COUNT; 22 createRate = DEFAULT_CREATE_RATE; 23 } 24 // Custom token bucket capacity and growth rate 25 public BucketUtil(int maxCount, int createRate) { 26 this.maxCount = maxCount; 27 this.createRate = createRate; 28 } 29 30 public int getSize() { 31 return size; 32 } 33 34 public boolean isFull() { 35 return size == maxCount; 36 } 37 38 //Generate a token according to the rate self increment 39 public synchronized void incrTokens() { 40 for (int i = 0; i < createRate; i++) 41 { 42 if (isFull()) 43 return; 44 size++; 45 } 46 } 47 48 // Take a token 49 public synchronized boolean getToken() { 50 if (size > 0) 51 size--; 52 else 53 return false; 54 return true; 55 } 56 57 @Override 58 public boolean equals(Object obj) { 59 if (obj == null) 60 return false; 61 BucketUtil bucket = (BucketUtil) obj; 62 if (bucket.size != size || bucket.createRate != createRate || bucket.maxCount != maxCount) 63 return false; 64 return true; 65 } 66 67 @Override 68 public int hashCode() { 69 return Objects.hash(maxCount, size, createRate); 70 } 71 72 }
Initialize token bucket
Initializes and generates a timer on the startup class
1 @EnableScheduling 2 @SpringBootApplication 3 public class DemoApplication { 4 5 public static void main(String[] args) { 6 SpringApplication.run(DemoApplication.class, args); 7 // For the convenience of testing, 1 capacity 1 growth rate is defined here 8 BucketUtil bucketUtil = new BucketUtil(1,1); 9 // The build name is: bucket Token bucket 10 BucketUtil.buckets.put("bucket",bucketUtil); 11 } 12 @Scheduled(fixedRate = 1000)// Timing 1 s 13 public void timer() { 14 if (BucketUtil.buckets.containsKey("bucket")){ 15 //Named: bucket The token bucket began to generate tokens continuously 16 BucketUtil.buckets.get("bucket").incrTokens(); 17 } 18 } 19 }
Custom annotations and exceptions
@Target({ElementType.METHOD})// METHOD Representative is used in method
@Retention(RetentionPolicy.RUNTIME)
public @interface BucketAnnotation {
}
1 public class APIException extends RuntimeException { 2 private static final long serialVersionUID = 1L; 3 private String msg; 4 public APIException(String msg) { 5 super(msg); 6 this.msg = msg; 7 } 8 }
Configuring Interceptors
1 /** 2 * Token bucket interceptor 3 */ 4 public class BucketInterceptor implements HandlerInterceptor { 5 6 // Preprocessing callback method, used before interface call true Representative release false Representative does not release 7 @Override 8 public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) { 9 if (!(handler instanceof HandlerMethod)) { 10 return true; 11 } 12 13 HandlerMethod handlerMethod = (HandlerMethod) handler; 14 Method method = handlerMethod.getMethod(); 15 16 BucketAnnotation methodAnnotation = method.getAnnotation(BucketAnnotation.class); 17 if (methodAnnotation!=null){ 18 // Under the name: bucket Get the token from the token bucket and release it when it is received. Throw an exception when it is not received 19 if(BucketUtil.buckets.get("bucket").getToken()){ 20 return true; 21 } 22 else{ 23 // Throw custom exception 24 throw new APIException("Sorry, you are restricted"); 25 } 26 }else { 27 return true; 28 } 29 } 30 // After the interface call, return to the previous use 31 @Override 32 public void postHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, ModelAndView modelAndView) throws Exception { 33 } 34 35 // After the entire request is completed, it is used before the view is rendered 36 @Override 37 public void afterCompletion(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, Exception e) throws Exception { 38 } 39 }
Inject interceptor
1 @Configuration 2 public class WebMvcConfg implements WebMvcConfigurer { 3 4 @Override 5 public void addInterceptors(InterceptorRegistry registry) { 6 // Token bucket interceptor add interceptor and select interception path 7 registry.addInterceptor(bucketInterceptor()).addPathPatterns("/**"); 8 } 9 @Bean 10 public BucketInterceptor bucketInterceptor() { 11 return new BucketInterceptor(); 12 } 13 }
AOP global exception handling
1 @RestControllerAdvice 2 public class WebExceptionControl { 3 @ExceptionHandler(APIException.class) 4 public E3Result APIExceptionHandler(APIException e) { 5 return E3Result.build(400,e.getMessage()); 6 } 7 }
test
Mark the user-defined annotation on the interface where we need to limit the current, as follows
@BucketAnnotation @RequestMapping(value = "/bucket") public E3Result bucket(){ return E3Result.ok("Access successful"); }
E3Result is only an encapsulated return class, which will not be posted here. Some of them can be replaced with their own, and others can be tested directly with String type
In order to facilitate the test, the capacity of the token bucket is set to 1, so this is the key to successful token retrieval
summary
The current limit above is only a demo, and there are many deficiencies, such as:
- Not applicable in distributed environment
- There can be multiple token buckets. When different interfaces adopt different token buckets, the interceptor cannot separate the current limit
- A request consumes one token at a time, which can be consumed maliciously
Improvement method:
- redis cluster access is adopted for token bucket
- Adding the value parameter to the annotation can mark the corresponding token bucket parameter to the corresponding interface. The interceptor needs to verify the annotation parameter to limit the current of multiple token buckets of multiple interfaces
- Limit the number of user IP verification to prevent malicious attacks
The current limit of the actual project will be more rigorous. The above just provides an idea and demonstration demo. Don't spray if you don't like it. Thank you.