Spring Boot 中使用 Redis + Aop 进行限流

news2024/9/30 21:23:42

Spring Boot 中使用 Redis 进行限流，通常你可以采用如下几种方式：

令牌桶算法（Token Bucket）
漏桶算法（Leaky Bucket）
固定窗口计数器（Fixed Window Counter）
滑动日志窗口（Sliding Log Window）

实现 Redis 限流，可以采用 Redis 提供的数据结构和功能脚本，如 Lua 脚本、Redisson 库等。以下是使用 Redis 和 Lua 脚本来实现令牌桶限流算法的示例：

步骤一：编写 Lua 脚本。

下面是一个限流的 Lua 脚本示例，实现基本的限流功能，放在Spring Boot项目下的resources目录下。

--获取KEY
local key = KEYS[1] -- 限流的 key

local limit = tonumber(ARGV[1]) --注解标注的限流次数

local curentLimit = tonumber(redis.call('get', key) or "0")

if curentLimit + 1 > limit
then return 0
else
    -- 自增长 1
    redis.call('INCRBY', key, 1)
    -- 设置过期时间
    redis.call('EXPIRE', key, ARGV[2])
    return curentLimit + 1
end

步骤二：定义限流注解

package your.package;

import java.lang.annotation.*;

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface RedisLimit {
    /**
     * 资源的key,唯一
     * 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 最多的访问限制次数
     */
    long permitsPerSecond() default 2;

    /**
     * 过期时间也可以理解为单位时间，单位秒，默认60
     */
    long expire() default 60;


    /**
     * 得不到令牌的提示语
     */
    String msg() default "系统繁忙,请稍后再试.";
}

步骤三：定义Aop切面类

package your.package;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

/**
 * Limit AOP
 */
@Slf4j
@Aspect
@Component
public class RedisLimitAop {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;


    @Pointcut("@annotation(your.package.RedisLimit)")
    private void check() {

    }

    private DefaultRedisScript<Long> redisScript;

    @PostConstruct
    public void init() {
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
    }


    @Before("check()")
    public void before(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        //拿到RedisLimit注解，如果存在则说明需要限流
        RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);

        if (redisLimit != null) {
            //获取redis的key
            String key = redisLimit.key();
            String className = method.getDeclaringClass().getName();
            String name = method.getName();

            String limitKey = key + className + method.getName();

            log.info(limitKey);

            if (StringUtils.isEmpty(key)) {
                throw new RedisLimitException("key cannot be null");
            }

            long limit = redisLimit.permitsPerSecond();

            long expire = redisLimit.expire();

            List<String> keys = new ArrayList<>();
            keys.add(key);

            Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));

            log.info("Access try count is {} for key={}", count, key);

            if (count != null && count == 0) {
                log.debug("获取key失败，key为{}", key);
                throw new RedisLimitException(redisLimit.msg());
            }
        }

    }

}

步骤四：自定义Redis限流异常

package your.package;


/**
 * Redis限流自定义异常
 * @date 2023/3/10 21:43
 */
public class RedisLimitException extends RuntimeException{
 public RedisLimitException(String msg) {
  super( msg );
 }
}

步骤五：自定义ResultInfo返回实体

package your.package;


import lombok.Getter;
import lombok.Setter;

@Getter
@Setter
public class ResultInfo<T> {

    private String message;
    private String code;
    private T data;


    public ResultInfo(String message, String code, T data) {
        this.message = message;
        this.code = code;
        this.data = data;
    }

    public static ResultInfo error(String message) {
        return new ResultInfo(message,"502",null);
    }



}

步骤六：定义Controller接口

package your.package;


import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@Slf4j
@RestController
@RequestMapping("/limit/redis")
public class LimitRedisController {

    /**
     * 基于Redis AOP限流
     */
    @GetMapping("/test")
    @RedisLimit(key = "redis-limit:test", permitsPerSecond = 2, expire = 1, msg = "当前排队人数较多，请稍后再试！")
    public String test() {
        log.info("限流成功。。。");
        return "ok";
    }

}