基于Redis+Lua实现分布式限流

1、新建一个Mavne项目，取名为rate_limiter，并引入Lombok和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、在rate_limiter项目下新建一个名为ratelimiter_annotation的子模块，在该模块的pom文件中添加redis的依赖。

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

3、在ratelimiter_annotation模块的src/main/java目录下建立service包，在service包下建立一个名为AccessLimiter的类。

@Service
@Slf4j
public class AccessLimiter {
    @Autowired
    private StringRedisTemplate redisTemplate;
    /**
     * DefaultRedisScript类用来加载脚本的，并设置相应的数据类型来接收lua脚本返回的数据，
     * 这个泛型类在使用时设置泛型是什么类型，脚本返回的结果就是用什么类型接收。
     * 该类只接收4种类型的返回类型(Long, Boolean, List, or deserialized value type)
    */
    @Autowired
    private DefaultRedisScript<Boolean> rateLimiterLua;

    public void limitAccess(String key,Integer limit){
        //执行lua脚本
        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、新建config包并建立名为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();
        //设置lua脚本
        redisScript.setLocation(new ClassPathResource("ratelimiter.lua"));
        //设置返回类型
        redisScript.setResultType(java.lang.Boolean.class);
        return redisScript;
    }
}

5、在resources目录下新建lua脚本文件ratelimiter.lua。

--
-- Created by IntelliJ IDEA.
-- User: wanglei
--
-- 在lua脚本中，有两个全局的变量，是用来接收redis应用端传递的键值和其它参数的，
-- 分别为KEYS、ARGV。

-- 在应用端传递给KEYS时是一个数组列表，在lua脚本中经过索引方式获取数组内的值。

-- 在应用端，传递给ARGV的参数比较灵活，能够是多个独立的参数，但对应到Lua脚本中是，
-- 统一用ARGV这个数组接收，获取方式也是经过数组下标获取。


-- 经过KEYS获取方法签名特征
local methodKey = KEYS[1]
redis.log(redis.LOG_DEBUG, 'key is', methodKey)

-- 经过ARGV传入限流大小
local limit = tonumber(ARGV[1])

-- 获取当前流量大小
local count = tonumber(redis.call('get', methodKey) or "0")

-- 是否超出限流阈值
if count + 1 > limit then
    -- 拒绝服务访问
    return false
else
    -- 没有超过阈值
    -- 设置当前访问的数量+1
    redis.call("INCRBY", methodKey, 1)
    -- 设置过时时间
    redis.call("EXPIRE", methodKey, 1)
    -- 放行
    return true
end

6、在rate_limiter项目中再新建一个ratelimiter_test的子模块用于测试咱们前面的脚本。在ratelimiter_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、在ratelimiter_test的src/main/java下新建controller包，并在controller包下建立一个TestController的类。

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

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

8、在application.properties中添加redis的配置

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

9、建立一个启动类并启动项目，在postman中测试一下查看限流的结果。

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

}

10、经过以上的几个步骤，已经实现了基于Redis+Lua的限流，可是代码还不够完美，如今咱们将项目改造一下，经过自定义的注解在项目的任何位置均可以实现限流。

先在ratelimiter_annotation模块中引入aop的依赖。

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

而后在ratelimiter_annotation模块中新建一个annotation的包，并在annotation包下建立一个名为AccessLimiter的注解。

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

再建立一个aspect的包，并建立一个名为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){
        //一、得到方法签名，做为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();

        //若是没有设置methodKey,从调用方法签名自动生成一个
        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;
            }
        }

        //二、调用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");
        }
    }
}

如今咱们九可使用咱们自定义的注解了，咱们在TestController新增一个方法

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

经过启动类再次启动咱们的项目并测试一下testAnnotation接口。