Rate limiting Spring-boot APIs using bucket4j and Spring MVC
Spring boot applications by default don’t have a mechanism to rate-limit the API calls that can be served to a client. Here is a simple way of implementing the token-bucket algorithm to rate limit our spring boot APIs.
Algorithm:
Token bucket algorithm uses the idea of a fixed “bucket” of possible requests that a client can make in a given time period. When a request is made, the bucket is checked and a token is given to the client. If there are tokens available, then the request is fulfilled.
Implementation:
To implement this technique, we will use a java library called “bucket4j”. Bucket4j is a java rate-limiting library that is based on the above algorithm. We can use this library to define “Buckets” and “Bandwidth” for our rest APIs.
High-level design:
Low-level design:
Adding the dependency :
We add the following dependency in the pom.xml file of the project to import the buckt4j library.
<dependency>
<groupId>com.github.vladimir-bukhtoyarov</groupId>
<artifactId>bucket4j-core</artifactId>
<version>6.0.1</version>
</dependency>Rate limiter service:
We create a rate limiter service that has a function to return a token bucket based on the API key that is provided.
Service:
public interface RateLimiterService {
Bucket resolveBucket(String apiKey);
}Service implementation:@Service
public class RateLimiterServiceImpl implements RateLimiterService {//Cache for API key and token bucket ( we can use any other caching method as well )
Map<String, Bucket> bucketCache = new ConcurrentHashMap<>();@Override
public Bucket resolveBucket(String apiKey) {
return bucketCache.computeIfAbsent(apiKey, this::newBucket);
}private Bucket newBucket(String s) {
return Bucket4j.builder()
.addLimit(Bandwidth.classic(10, Refill.intervally(10, Duration.ofMinutes(1)))).build();
}}
Rate Limit Interceptor:
We will create a custom HandlerInterceptor (spring framework servlet class) and call it RateLimitInterceptor. This class will be used to override the “preHandle” method which is used to handle all the incoming API requests to the spring-boot application.
We’ll autowire the RateLimiterService to create and assign a token bucket for every valid API key sent in the requests.
Using a ConsumptionProbe instance, that checks how many tokens are left, we can inform the client about the status of the API requests that remain and if the limit is exhausted, how long it would take to replenish the token bucket.
If we have enough tokens left for the API key, we let the request pass else we return a 426 HTTP code to the client.
@Component
public class RateLimitInterceptor implements HandlerInterceptor {@Autowired
private RateLimiterService rateLimiterService;@Override
public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
throws Exception {
String apiKey = request.getHeader(“api-key”);
if (apiKey == null || apiKey.isEmpty()) {
response.sendError(HttpStatus.BAD_REQUEST.value(), “Missing Header: api-key”);
return false;
}Bucket bucket = rateLimiterService.resolveBucket(apiKey);
ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
if(probe.isConsumed()) {
response.addHeader(“X-Rate-Limit-Remaining”, String.valueOf(probe.getRemainingTokens()));
return true;
} else {
long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
response.addHeader(“Rate-Limit-Retry-After-Seconds”, String.valueOf(waitForRefill));
response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(),
“You have exhausted your API Request Quota”);
return false;
}
}
}
AppConfig:
Finally, we add the custom interceptor to our WebMvcConfigurer and add it to the spring interceptor registry using a config class.
We can add the API endpoint patterns for which our custom interceptor is to be used. If we want it for all the APIs in the application, we can simply say “/**”.
@Configuration
public class AppConfig implements WebMvcConfigurer {
@Autowired
private RateLimitInterceptor rateLimitInterceptor;
@Override
public void addInterceptors(InterceptorRegistry registry) {
registry.addInterceptor(rateLimitInterceptor).addPathPatterns(“/**”);
}
}
