How to Implement Rate Limiting in Spring Boot with Custom Annotations and Redis

Implement request rate limiting so that a specific interface can be called only N times within a defined time period.

Principle

The server records each request count in Redis; if the count exceeds the limit, the request is denied. The Redis key has an expiration time, so it is automatically removed after the period.

Code Implementation

To give the solution a more professional feel, it is implemented using a custom annotation.

@RequestLimit Annotation

import java.lang.annotation.*;
/**
 * Custom annotation for request limiting
 */
@Documented
@Inherited
@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestLimit {
    // Maximum requests allowed within 'second' seconds
    int second() default 1;
    int maxCount() default 1;
}

RequestLimitIntercept Interceptor

A custom interceptor checks the request count before processing.

import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;
import top.lrshuai.limit.annotation.RequestLimit;
import top.lrshuai.limit.common.ApiResultEnum;
import top.lrshuai.limit.common.Result;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;
/**
 * Request interceptor
 */
@Slf4j
@Component
public class RequestLimitIntercept extends HandlerInterceptorAdapter {
    @Autowired
    private RedisTemplate<String,Object> redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler.getClass().isAssignableFrom(HandlerMethod.class)) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            RequestLimit methodAnnotation = method.getAnnotation(RequestLimit.class);
            RequestLimit classAnnotation = method.getDeclaringClass().getAnnotation(RequestLimit.class);
            RequestLimit requestLimit = methodAnnotation != null ? methodAnnotation : classAnnotation;
            if (requestLimit != null) {
                if (isLimit(request, requestLimit)) {
                    resonseOut(response, Result.error(ApiResultEnum.REQUST_LIMIT));
                    return false;
                }
            }
        }
        return super.preHandle(request, response, handler);
    }

    // Determine if request is limited
    public boolean isLimit(HttpServletRequest request, RequestLimit requestLimit) {
        String limitKey = request.getServletPath() + request.getSession().getId();
        Integer redisCount = (Integer) redisTemplate.opsForValue().get(limitKey);
        if (redisCount == null) {
            // initial count
            redisTemplate.opsForValue().set(limitKey, 1, requestLimit.second(), TimeUnit.SECONDS);
        } else {
            if (redisCount.intValue() >= requestLimit.maxCount()) {
                return true;
            }
            // increment count
            redisTemplate.opsForValue().increment(limitKey);
        }
        return false;
    }

    /**
     * Write response back to client
     */
    private void resonseOut(HttpServletResponse response, Result result) throws IOException {
        response.setCharacterEncoding("UTF-8");
        response.setContentType("application/json; charset=utf-8");
        String json = JSONObject.toJSON(result).toString();
        PrintWriter out = response.getWriter();
        out.append(json);
    }
}

After writing the interceptor, it must be registered.

WebMvcConfig Configuration Class

For Spring Boot 2.x, implement WebMvcConfigurer; for Spring Boot 1.x, extend WebMvcConfigurerAdapter. Override addInterceptors() to add the custom interceptor.

@Slf4j
@Component
public class WebMvcConfig implements WebMvcConfigurer {
    @Autowired
    private RequestLimitIntercept requestLimitIntercept;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        log.info("Adding request limit interceptor");
        registry.addInterceptor(requestLimitIntercept);
    }
}

Controller Example

Test endpoints with the annotation.

Usage

Apply on a class: @RequestLimit(maxCount = 5, second = 1) Apply on a method: @RequestLimit(maxCount = 5, second = 1) Parameters: maxCount – maximum number of requests; second – time window in seconds.

Default: one request per second per interface.

@RestController
@RequestMapping("/index")
@RequestLimit(maxCount = 5, second = 1)
public class IndexController {
    /**
     * @RequestLimit on method takes precedence over class parameters
     */
    @GetMapping("/test1")
    @RequestLimit
    public Result test(){
        // TODO ...
        return Result.ok();
    }

    /**
     * @RequestLimit on class uses class parameters
     */
    @GetMapping("/test2")
    public Result test2(){
        // TODO ...
        return Result.ok();
    }
}

If both class and method have @RequestLimit, the method's parameters are used.

Code Repository

Full source code is available at:

Gitee: https://gitee.com/rstyro/spring-boot/tree/master/SpringBoot-limit

GitHub: https://github.com/rstyro/Springboot/tree/master/SpringBoot-limit

If you find the example useful, please give it a star on GitHub.