Redis does interface current limiting, which is an annotation!

In addition to caching, Redis can also do many things: distributed locking, flow limiting, idempotency of request processing interface... Too much, too much ~

Today, I want to talk with my friends about how to use Redis to deal with interface current limitation. This is also the knowledge point involved in the recent TienChin project. I'll bring it out to talk about this topic with you, and I'll talk about it in the video later.

1. Preparation

First, we create a Spring Boot project and introduce Web and Redis dependencies. At the same time, considering that the interface current limit is generally marked through annotations, and the annotations are resolved through AOP, we also need to add AOP dependencies. The final dependencies are as follows:

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

Then prepare a Redis instance in advance. After the project is configured, we can directly configure the basic information of Redis, as follows:

spring.redis.host=localhost
spring.redis.port=6379
spring.redis.password=123

Well, the preparations are in place.

2. Current limiting notes

Next, we create a current limiting annotation. We divide current limiting into two cases:

  1. Global current limiting for the current interface. For example, the interface can be accessed 100 times in one minute.
  2. Current limiting for an IP address. For example, an IP address can be accessed 100 times in one minute.

For these two cases, we create an enumeration class:

public enum LimitType {
    /**
     * Default policy global current limit
     */
    DEFAULT,
    /**
     * Current limiting according to the requester's IP
     */
    IP
}

Next, let's create a current limiting annotation:

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter {
    /**
     * Current limiting key
     */
    String key() default "rate_limit:";

    /**
     * Current limiting time in seconds
     */
    int time() default 60;

    /**
     * Current limiting times
     */
    int count() default 100;

    /**
     * Current limiting type
     */
    LimitType limitType() default LimitType.DEFAULT;
}

The first parameter is the current limiting key, which is only a prefix. In the future, the complete key is this prefix plus the complete path of the interface method to form the current limiting key, which will be stored in Redis.

The other three parameters are easy to understand, so I won't say more.

Well, in the future, add @ RateLimiter annotation to any interface that needs current limiting, and then configure relevant parameters.

3. Customize RedisTemplate

My friends know that in Spring Boot, we are more used to using Spring Data Redis to operate Redis, but the default RedisTemplate has a small hole, that is, JdkSerializationRedisSerializer is used for serialization. I don't know if my friends have noticed that the key and value stored on Redis directly with this serialization tool will have more prefixes in the future, This leads to errors when you read with the command.

For example, when storing, the key is name and the value is javaboy, but when you operate on the command line, get name can't get the data you want. The reason is that there are more characters in front of the name after saving to redis. At this time, you can only continue to use RedisTemplate to read and retrieve it.

The Lua script will be used when we use Redis for current limiting. When using Lua script, the above situation will occur, so we need to modify the serialization scheme of RedisTemplate.

Some friends may say why not use StringRedisTemplate? StringRedisTemplate does not have the problems mentioned above, but the data types it can store are not rich enough, so it is not considered here.

Modify the RedisTemplate serialization scheme. The code is as follows:

@Configuration
public class RedisConfig {

    @Bean
    public RedisTemplate<Object, Object> redisTemplate(RedisConnectionFactory connectionFactory) {
        RedisTemplate<Object, Object> redisTemplate = new RedisTemplate<>();
        redisTemplate.setConnectionFactory(connectionFactory);
        // Replace the default serialization with Jackson2JsonRedisSerialize (JDK serialization is adopted by default)
        Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<>(Object.class);
        ObjectMapper om = new ObjectMapper();
        om.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        om.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        jackson2JsonRedisSerializer.setObjectMapper(om);
        redisTemplate.setKeySerializer(jackson2JsonRedisSerializer);
        redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashKeySerializer(jackson2JsonRedisSerializer);
        redisTemplate.setHashValueSerializer(jackson2JsonRedisSerializer);
        return redisTemplate;
    }
}

There's nothing to say about this. We use the default jackson serialization method in Spring Boot to solve both key and value.

4. Develop Lua script

In fact, I mentioned in the previous vhr video that some atomic operations in Redis can be realized with the help of lua script. We have two different ideas to call Lua script:

  1. Define the Lua script on the Redis server, and then calculate a hash value. In the Java code, use this hash value to lock which Lua script to execute.
  2. Directly define the Lua script in Java code and send it to Redis server for execution.

Spring Data Redis also provides an interface to operate Lua scripts, which is relatively convenient, so we adopt the second scheme here.

We create a new lua folder in the resources directory to store lua scripts. The script contents are as follows:

local key = KEYS[1]
local count = tonumber(ARGV[1])
local time = 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)

This script is actually not difficult. You probably know what to do at a glance. KEYS and ARGV are parameters passed in when calling later. tonumber is to convert a string into a number, redis Call is to execute specific redis instructions. The specific process is as follows:

  1. First, get the key passed in and the count and time of current limit.
  2. Get the value corresponding to this key through get. This value is how many times this interface can be accessed in the current time window.
  3. If it is the first visit, the result obtained at this time is nil, otherwise the result obtained should be a number, so judge next. If the result obtained is a number and this number is greater than count, it means that the traffic limit has been exceeded, and then directly return the query result.
  4. If the result obtained is nil, it indicates that it is the first access. At this time, the current key will be incremented by 1, and then an expiration time will be set.
  5. Finally, return the value after increasing by 1.

In fact, this Lua script is very easy to understand.

Next, we load the Lua script in a Bean, as follows:

@Bean
public DefaultRedisScript<Long> limitScript() {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
    redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
    redisScript.setResultType(Long.class);
    return redisScript;
}

OK, our Lua script is ready now.

5. Annotation analysis

Next, we need to customize the section to resolve this annotation. Let's take a look at the definition of the section:

@Aspect
@Component
public class RateLimiterAspect {
    private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);

    @Autowired
    private RedisTemplate<Object, Object> redisTemplate;

    @Autowired
    private RedisScript<Long> limitScript;

    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable {
        String key = rateLimiter.key();
        int time = rateLimiter.time();
        int count = rateLimiter.count();

        String combineKey = getCombineKey(rateLimiter, point);
        List<Object> keys = Collections.singletonList(combineKey);
        try {
            Long number = redisTemplate.execute(limitScript, keys, count, time);
            if (number==null || number.intValue() > count) {
                throw new ServiceException("Access is too frequent. Please try again later");
            }
            log.info("Restriction request'{}',Current request'{}',cache key'{}'", count, number.intValue(), key);
        } catch (ServiceException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException("The flow limit of the server is abnormal. Please try again later");
        }
    }

    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP) {
            stringBuffer.append(IpUtils.getIpAddr(((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest())).append("-");
        }
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
        return stringBuffer.toString();
    }
}

The annotation of @ rater is added to the pre annotation of all aspects of the notification.

  1. First, get the key, time and count parameters in the annotation.
  2. Get a combined key. The so-called combined key is to add the full path of the method based on the annotated key attribute. If it is in IP mode, add the IP address. Taking IP mode as an example, the final generated key is similar to this: rate_ limit:127.0.0.1-org. javaboy. ratelimiter. controller. Hellocontroller hello (if it is not in the IP mode, the generated key does not contain the IP address).
  3. Put the generated key into the collection.
  4. Through redistemplate The execute method executes a Lua script. The first parameter is the object encapsulated by the script, and the second parameter is key, which corresponds to the KEYS in the script, followed by the variable length parameter, which corresponds to the ARGV in the script.
  5. Compare the execution result of Lua script with count. If it is greater than count, it means that it is overloaded and throw exceptions.

All right, it's done.

6. Interface test

Next, we will conduct a simple test of the interface, as follows:

@RestController
public class HelloController {
    @GetMapping("/hello")
    @RateLimiter(time = 5,count = 3,limitType = LimitType.IP)
    public String hello() {
        return "hello>>>"+new Date();
    }
}

Each IP address can only be accessed 3 times in 5 seconds.

This can be tested by manually refreshing the browser.

7. Global exception handling

Since exceptions are thrown when overloaded, we also need a global exception handler, as follows:

@RestControllerAdvice
public class GlobalException {
    @ExceptionHandler(ServiceException.class)
    public Map<String,Object> serviceException(ServiceException e) {
        HashMap<String, Object> map = new HashMap<>();
        map.put("status", 500);
        map.put("message", e.getMessage());
        return map;
    }
}

This is a small demo. Instead of defining entity classes, I will directly use Map to return JSON.

All right, it's done.

Finally, let's look at the test effect under overload:

Well, this is how we use Redis to limit current.

SongGe is currently recording the TienChin project video, and the content of this article is also a part of the TienChin project video ~ the TienChin project adopts the Spring Boot+Vue3 technology stack, which will involve all kinds of fun technologies. My friends come to work with SongGe on a project with a completion rate of more than 90%.

Tags: Redis Spring Spring Boot Spring MVC

Posted by harley1387 on Tue, 17 May 2022 16:48:56 +0300