Distributed current limiting based on redis + Lua

Time:2020-10-1

1、 Create a mavne project named rate_ Limiter and introduce the dependency of Lombok and guava.

<dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
</dependency>
<dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
    <version>29.0-jre</version>
</dependency>

2、 At rate_ Under the limiter project, create a new one named ratelimiter_ Annotation sub module, add redis dependency in the POM file of the module.

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

3、 At ratelimiter_ Create the service package in the Src / main / Java directory of the annotation module, and create a class named accesslimiter under the service package.

@Service
@Slf4j
public class AccessLimiter {
    @Autowired
    private StringRedisTemplate redisTemplate;
    /**
     *The defaultredisscript class is used to load the script, and set the corresponding data type to receive the data returned by Lua script,
     *This generic class sets the type of generics when it is used, and the result returned by the script is what type to receive.
     *This class only receives four types of return types (long, Boolean, list, or deserialized value type)
    */
    @Autowired
    private DefaultRedisScript<Boolean> rateLimiterLua;

    public void limitAccess(String key,Integer limit){
        //Execute Lua script
        boolean acquire=redisTemplate.execute(
                rateLimiterLua,
                Lists.newArrayList(key),
                limit.toString());
        if (!acquire){
            log.error("your access is blocked,key={}",key);
            throw new RuntimeException("your access is blocked");
        }
    }
}

4、 Create a new config package and create a configuration class named redisconfiguration

@Configuration
public class RedisConfiguration {

    @Bean
    public RedisTemplate<String,String> redisTemplate(
            RedisConnectionFactory factory
    ){
        return new StringRedisTemplate(factory);
    }

    @Bean
    public DefaultRedisScript loadRedisScript(){
        DefaultRedisScript redisScript=new DefaultRedisScript();
        //Set Lua script
        redisScript.setLocation(new ClassPathResource("ratelimiter.lua"));
        //Set return type
        redisScript.setResultType(java.lang.Boolean.class);
        return redisScript;
    }
}

5、 Create a Lua script file in the resources directory ratelimiter.lua 。

--
-- Created by IntelliJ IDEA.
-- User: wanglei
--
--In Lua script, there are two global variables that are used to receive key values and other parameters passed by the redis application,
--They were keys and argv.

--When passed to keys on the application side, it is an array list. In Lua script, the values in the array are obtained by index.

--On the application side, the parameters passed to argv are more flexible. They can be multiple independent parameters, but the corresponding parameters in Lua script are,
--Argv is used to receive the array uniformly, and the acquisition method is also obtained through the array subscript.


--Obtaining method signature features through keys
local methodKey = KEYS[1]
redis.log(redis.LOG_DEBUG, 'key is', methodKey)

--The current limiting size is transmitted through argv
local limit = tonumber(ARGV[1])

--Get current traffic size
local count = tonumber(redis.call('get', methodKey) or "0")

--Is the current limiting threshold exceeded
if count + 1 > limit then
    --Denial of service access
    return false
else
    --The threshold was not exceeded
    --Set number of current visits + 1
    redis.call("INCRBY", methodKey, 1)
    --Set expiration time
    redis.call("EXPIRE", methodKey, 1)
    --Release
    return true
end

6、 At rate_ Create a ratelimiter in the limiter project_ The sub module of test is used to test our previous script. At ratelimiter_ The following dependencies are introduced in test.

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

<dependency>
    <groupId>${project.groupId}</groupId>
    <artifactId>ratelimiter_annotation</artifactId>
    <version>${project.version}</version>
</dependency>

7、 At ratelimiter_ Create a new controller package under Src / main / Java of test, and create a class of testcontroller under the controller package.

@RestController
@Slf4j
public class TestController {
    @Autowired
    private AccessLimiter accessLimiter;

    @GetMapping("test")
    public String test(){
        accessLimiter.limitAccess("ratelimiter-test",1);
        return "success";
    }
}

8、 In application.properties Add redis configuration in

spring.redis.database=0
spring.redis.host=localhsot
spring.redis.port=6379
spring.redis.password=root

9、 Create a startup class and start the project. Test it in postman to see the result of current restriction.

@SpringBootApplication
public class RatelimiterTestApplication {
    
    public static void main(String[] args) {
        SpringApplication.run(RatelimiterTestApplication.class, args);
    }

}

10、 Through the above steps, we have realized the current restriction based on redis + Lua, but the code is not perfect. Now we will transform the project, and we can realize the current restriction at any position of the project through the custom annotation.

First at ratelimiter_ AOP dependency is introduced into annotation module.

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

And then at ratelimiter_ In the annotation module, create a new annotation package and an annotation named accesslimiter under the annotation package.

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface AccessLimiter {
    int limit();
    String methodKey() default "";
}

Create an aspect package, and create a class called accesslimiteraspect

@Slf4j
@Aspect
@Component
public class AccessLimiterAspect {

    @Autowired
    private StringRedisTemplate redisTemplate;
    @Autowired
    private DefaultRedisScript<Boolean> rateLimiterLua;

    @Pointcut("@annotation(com.wl.annotation.AccessLimiter)")
    public void cut(){
        log.info("cut");
    }

    @Before("cut()")
    public void before(JoinPoint joinPoint){
        //1. Get the method signature as the method key
        MethodSignature methodSignature= (MethodSignature) joinPoint.getSignature();
        Method method=methodSignature.getMethod();
        AccessLimiter annotation=method.getAnnotation(AccessLimiter.class);
        if (annotation==null){
            return;
        }
        String key=annotation.methodKey();
        Integer limit=annotation.limit();

        //If the methodkey is not set, it will automatically generate one from the calling method signature
        if (StringUtils.isEmpty(key)){
            Class[] type=method.getParameterTypes();
            key=method.getName();

            if (type!=null){
                String paramTypes= Arrays.stream(type)
                        .map(Class::getName)
                        .collect(Collectors.joining(","));
                log.info("param types: "+paramTypes);
                key+="#"+paramTypes;
            }
        }

        //2. Call redis
        boolean acquire=redisTemplate.execute(
                rateLimiterLua,
                Lists.newArrayList(key),
                limit.toString());
        if (!acquire){
            log.error("your access is blocked,key={}",key);
            throw new RuntimeException("your access is blocked");
        }
    }
}

Now that we can use our custom annotation, we add a method to testcontroller

@GetMapping("test-annotation")
@com.wl.annotation.AccessLimiter(limit = 1)
public String testAnnotation(){
    return "success";
}

Start our project again with the startup class and test the testannotation interface.