Implementing Token Bucket Rate Limiting in Spring Boot with Custom Annotation and Interceptor

In high‑concurrency scenarios, rate limiting is a common backend technique; this guide shows a complete demo that combines the token‑bucket algorithm, a Spring interceptor, custom annotations, and custom exceptions to enforce limits on system, interface, or user access.

Token Bucket Concept

A fixed‑size bucket continuously generates tokens at a constant rate. If tokens are not consumed, they accumulate until the bucket is full; excess tokens overflow. Each request must take a token; if none are available, the request is rejected as rate‑limited.

Implementation Steps

Create a token‑bucket class.

Initialize the bucket at application startup and schedule periodic token addition.

Define a custom rate‑limit annotation to mark protected endpoints.

Configure an interceptor that checks the annotation, attempts to acquire a token, and throws a custom exception when the bucket is empty.

Use AOP to handle the custom exception globally.

Bucket Utility Class

The BucketUtil class manages token count, capacity, and generation rate, with synchronized methods to ensure thread safety.

public class BucketUtil {
    static final int DEFAULT_MAX_COUNT = 10;
    static final int DEFAULT_CREATE_RATE = 1;
    public static final HashMap<String, BucketUtil> buckets = new HashMap<>(10);
    final int maxCount;
    int createRate;
    int size = 0;
    public BucketUtil() { maxCount = DEFAULT_MAX_COUNT; createRate = DEFAULT_CREATE_RATE; }
    public BucketUtil(int maxCount, int createRate) { this.maxCount = maxCount; this.createRate = createRate; }
    public synchronized void incrTokens() { for (int i = 0; i < createRate; i++) { if (size == maxCount) return; size++; } }
    public synchronized boolean getToken() { if (size > 0) { size--; return true; } return false; }
    // equals and hashCode omitted for brevity
}

Initializing the Bucket

In the main application class, a bucket is created and a scheduled task adds tokens every second.

@EnableScheduling
@SpringBootApplication
public class DemoApplication {
    public static void main(String[] args) {
        SpringApplication.run(DemoApplication.class, args);
        BucketUtil bucketUtil = new BucketUtil(1, 1);
        BucketUtil.buckets.put("bucket", bucketUtil);
    }
    @Scheduled(fixedRate = 1000)
    public void timer() {
        if (BucketUtil.buckets.containsKey("bucket")) {
            BucketUtil.buckets.get("bucket").incrTokens();
        }
    }
}

Custom Annotation and Exception

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface BucketAnnotation {}

public class APIException extends RuntimeException {
    private static final long serialVersionUID = 1L;
    private String msg;
    public APIException(String msg) { super(msg); this.msg = msg; }
}

Interceptor Configuration

The BucketInterceptor checks for the annotation, attempts to acquire a token from the bucket named "bucket", and either allows the request or throws APIException.

public class BucketInterceptor implements HandlerInterceptor {
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        if (!(handler instanceof HandlerMethod)) return true;
        Method method = ((HandlerMethod) handler).getMethod();
        BucketAnnotation ann = method.getAnnotation(BucketAnnotation.class);
        if (ann != null) {
            if (BucketUtil.buckets.get("bucket").getToken()) return true;
            else throw new APIException("不好意思，您被限流了");
        }
        return true;
    }
    // postHandle and afterCompletion omitted
}

Registering the Interceptor

@Configuration
public class WebMvcConfg implements WebMvcConfigurer {
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(bucketInterceptor()).addPathPatterns("/**");
    }
    @Bean
    public BucketInterceptor bucketInterceptor() { return new BucketInterceptor(); }
}

Global Exception Handling with AOP

@RestControllerAdvice
public class WebExceptionControl {
    @ExceptionHandler(APIException.class)
    public E3Result APIExceptionHandler(APIException e) {
        return E3Result.build(400, e.getMessage());
    }
}

Testing the Rate Limit

Apply the custom annotation to any controller method to protect it.

@BucketAnnotation
@RequestMapping("/bucket")
public E3Result bucket() { return E3Result.ok("访问成功"); }

When the bucket capacity is set to 1, the first request succeeds (token acquired) and subsequent rapid requests are rejected, demonstrating the rate‑limiting behavior.

Summary and Improvements

The demo illustrates basic token‑bucket rate limiting but has limitations: it does not work in distributed environments, cannot handle multiple buckets per endpoint, and each request consumes a token which could be abused. Suggested enhancements include using Redis for a distributed bucket, adding a value attribute to the annotation for per‑endpoint bucket selection, and incorporating IP‑based throttling.