SpringBoot implements current limiting annotations

SpringBoot implements current limiting annotations

In a high-concurrency system, there are three ways to protect the system: caching, downgrading, and current limiting.

The purpose of current limiting is to protect the system by limiting the rate of concurrent access requests or the number of requests within a time window. Once the limit rate is reached, the service can be denied, queued or waited

1. Current limiting type enumeration class

/**
 * 限流类型
 * @author ss_419
 */

public enum LimitType {
    
    
    /**
     * 默认的限流策略，针对某一个接口进行限流
     */
    DEFAULT,
    /**
     * 针对某一个IP进行限流
     */
    IP

}

2. Custom current limit annotation

/**
 * @author ss_419
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimiter {
    
    
    /**
     * 限流的 key，主要是指前缀
     * @return
     */
    String key() default "rate_limit:";

    /**
     * 在时间窗内的限流次数
     * @return
     */
    int count() default 100;

    /**
     * 限流类型
     * @return
     */
    LimitType limitType() default LimitType.DEFAULT;
    /**
     * 限流时间窗
     * @return
     */
    int time() default 60;
}

3. Current limiting lua script

1. Since we use Redis for current limiting, we need to introduce the maven dependency of Redis and the dependency of aop at the same time

<!-- aop依赖 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>
<!-- redis依赖 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

2. Configure redis and lua scripts

@Configuration
public class RedisConfig {
    
    

    @Bean
    RedisTemplate<Object,Object> redisTemplate(RedisConnectionFactory factory) {
    
    
        RedisTemplate<Object, Object> template = new RedisTemplate<>();
        template.setConnectionFactory(factory);
        Jackson2JsonRedisSerializer jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer(Object.class);
        template.setKeySerializer(jackson2JsonRedisSerializer);
        template.setValueSerializer(jackson2JsonRedisSerializer);
        template.setHashKeySerializer(jackson2JsonRedisSerializer);
        template.setHashValueSerializer(jackson2JsonRedisSerializer);

        return template;
    }

    /**
     * 读取lua脚本
     * @return
     */
    @Bean
    DefaultRedisScript<Long> limitScript() {
    
    
        DefaultRedisScript<Long> script = new DefaultRedisScript<>();
        script.setResultType(Long.class);
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
        return script;
    }
}

Through the Lua script, according to the key value cached in Redis, it is judged whether the number of visits exceeds the number of visits within the current limit time (also the expiration time of the key). If the number of visits is not exceeded, the number of visits will be +1, and true will be returned. Otherwise, false will be returned.
limit.lua:

local key = KEYS[1]
local time = tonumber(ARGV[1])
local count = tonumber(ARGV[2])
local current = redis.call('get', key)
if current and tonumber(current) > count then
    return tonumber(current)
end
current = redis.call('incr', key)
if tonumber(current) == 1 then
    redis.call('expire', key, time)
end
return tonumber(current)

4. Current-limiting section processing class

1. Use our Lua script just now to judge whether the number of current limits has been exceeded, and return a custom exception after exceeding the number of current limits, and then catch the exception in the global exception and return JSON data.

2. According to the annotation parameters, determine the current limit type, and stitch the cache key value

package org.pp.ratelimiter.aspectj;


import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.pp.ratelimiter.annotation.RateLimiter;
import org.pp.ratelimiter.enums.LimitType;
import org.pp.ratelimiter.exception.RateLimitException;
import org.pp.ratelimiter.utils.IpUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Collections;


@Aspect
@Component
public class RateLimiterAspect {
    
    

    private static final Logger logger = LoggerFactory.getLogger(RateLimiterAspect.class);

    @Autowired
    RedisTemplate<Object, Object> redisTemplate;

    @Autowired
    RedisScript<Long> redisScript;


    @Before("@annotation(rateLimiter)")
    public void before(JoinPoint jp, RateLimiter rateLimiter) throws RateLimitException {
    
    
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        String combineKey = getCombineKey(rateLimiter, jp);
        try {
    
    
            Long number = redisTemplate.execute(redisScript, Collections.singletonList(combineKey), time, count);
            if (number == null || number.intValue() > count) {
    
    
                //超过限流阈值
                logger.info("当前接口以达到最大限流次数");
                throw new RateLimitException("访问过于频繁，请稍后访问");
            }
            logger.info("一个时间窗内请求次数：{}，当前请求次数：{}，缓存的 key 为 {}", count, number, combineKey);
        } catch (Exception e) {
    
    
            throw e;
        }
    }

    /**
     * 这个 key 其实就是接口调用次数缓存在 redis 的 key
     * rate_limit:11.11.11.11-org.javaboy.ratelimit.controller.HelloController-hello
     * rate_limit:org.javaboy.ratelimit.controller.HelloController-hello
     * @param rateLimiter
     * @param jp
     * @return
     */
    private String getCombineKey(RateLimiter rateLimiter, JoinPoint jp) {
    
    
        StringBuffer key = new StringBuffer(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP) {
    
    
            key.append(IpUtils.getIpAddr(((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest()))
                    .append("-");
        }
        MethodSignature signature = (MethodSignature) jp.getSignature();
        Method method = signature.getMethod();
        key.append(method.getDeclaringClass().getName())
                .append("-")
                .append(method.getName());
        return key.toString();
    }
}

5. Use and test

@RestController
public class HelloController {
    
    

    /**
     * 限流 10 秒之内，这个接口可以访问3次
     * @return
     */
    @GetMapping("/hello")
    @RateLimiter(time = 10,count = 3)
    public Map<String, Object> hello() {
    
    
        Map<String, Object> map = new HashMap<>();
        map.put("status", 200);
        map.put("message", "Hello RateLimiter");
        return map;
    }

}

An exception will be reported if the number of visits exceeds 3 within ten seconds

For the data in redis, add 1 for each access.

When the number of accesses exceeds 3, the current limit operation will be performed.