基於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接口。