基于Redis+AOP+Lua脚本实现一个服务器限流机制

后端某些接口在高并发的压力下往往会导致性能的严重下降，为了维持我们后端服务型的高性能和高可用，我们往往可以对某些接口或某些用户去设计限流机制，控制这些热点接口的访问量，我这里利用Redis的高性能优势，并整合AOP编程和引入Lua限流脚本在SpringBoot中对任意接口或某些用户实现了访问量限流的机制，其中，我这里给出了三种限流机制：用户，IP地址，全局限流

1.定义限流方式

 /**
 * 限流类型
 * @Author GuihaoLv
 */
public enum LimitType {
    /**
     * 默认策略全局限流
     */
    DEFAULT,

    /**
     * 根据请求者IP进行限流
     */
    IP,

    /**
     * 根据请求者的用户ID进行限流
     */
    USER,

    /**
     * 根据请求者的部门进行限流
     */
    DEPT,
}

2.引入AOP，自定义限流注解和限流处理的切面类

/**
 * 限流注解
 * @Author GuihaoLv
 */
//实例 @RateLimiter(time = 60, count = 5, limitType = LimitType.IP) 效果：同一IP 60秒内最多允许5次登录尝试。
@Target(ElementType.METHOD) //表示该注解仅能标注在方法上，用于对具体方法进行限流控制。
@Retention(RetentionPolicy.RUNTIME) //注解在运行时保留，可通过反射机制读取注解信息，实现动态限流逻辑。
@Documented //注解信息会包含在生成的 JavaDoc 中
public @interface RateLimiter {
    /**
     * 限流key
     */
    public String key() default RedisConstant.RATE_LIMIT_KEY;
    /**
     * 限流时间,单位秒
     */
    public int time() default 60;

    /**
     * 限流次数
     */
    public int count() default 100;

     /**
     * 限流类型
     */
    public LimitType limitType() default LimitType.DEFAULT;
}

/**
 * 限流处理切面
 * @Author GuihaoLv
 */
@Aspect
@Component
//确保仅当配置项 spring.cache.type=redis 时，切面才会生效
@ConditionalOnProperty(prefix = "spring.cache", name = { "type" }, havingValue = "redis", matchIfMissing = false)
public class RateLimiterAspect
{
    private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);

    private RedisTemplate<Object, Object> redisTemplate; //redis 操作模板，用于执行 Lua 脚本和 Redis 命令

    private RedisScript<Long> limitScript; //限流核心逻辑的 Lua 脚本

    @Autowired
    public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate)
    {
        this.redisTemplate = redisTemplate;
    }

    @Autowired
    public void setLimitScript(RedisScript<Long> limitScript)
    {
        this.limitScript = limitScript;
    }
    @Autowired
    private ObjectMapper jacksonObjectMapper;

    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable
    {
        //1. 获取注解参数
        int time = rateLimiter.time(); //时间窗口（秒）
        int count = rateLimiter.count();// 允许的请求次数

        //2. 生成唯一限流 Key
        String combineKey = getCombineKey(rateLimiter, point);
        List<Object> keys = Collections.singletonList(combineKey);
        try
        {
            //3. 以key为参数执行 Lua 脚本（原子性操作） keys:Redis 存储的 Key
            //Lua脚本会检查Key是否存在，如果不存在则创建并设置过期时间，如果存在则递增计数器。
            Long number = redisTemplate.execute(limitScript, keys, count, time);
            if (StringUtils.isEmpty(number) || number.intValue() > count)
            {
                throw new Exception("访问过于频繁，请稍候再试");
            }
            log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);
        } catch (RuntimeException e)
        {
            throw new RuntimeException("服务器限流异常，请稍候再试");
        } catch (Exception e)
        {
            throw e;
        }
    }


    //IP + 类名 + 方法名 的拼接方式，确保不同场景的Key不冲突。
    //Key结构清晰，便于调试和监控（如通过Redis直接查看计数器）。
    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point)
    {
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());

        switch (rateLimiter.limitType()) {
            case IP:
                try {
                    limitByIp(stringBuffer);
                } catch (IOException e) {
                    e.printStackTrace();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                break;
            case USER:
                limitByUser(stringBuffer);
                break;
            case DEFAULT:
                limitByDefault(stringBuffer,point);
                break;
        }

           return stringBuffer.toString();
    }


     /**
     * 按IP限流
     * @param stringBuffer
     */
    private final HttpClient httpClient = HttpClient.newHttpClient();
    private void limitByIp(StringBuffer stringBuffer) throws IOException, InterruptedException {
        String[] services = {
                "https://api.ipify.org",
                "https://icanhazip.com"
        };
        for (String serviceUrl : services) {
                HttpRequest request = HttpRequest.newBuilder()
                        .uri(URI.create(serviceUrl))
                        .GET()
                        .build();
                HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
                String ip = response.body().trim();
                stringBuffer.append(":ip:").append(ip);
        }
    }

    /**
     * 全局限流
     * @param stringBuffer
     * @param point
     */
    private void limitByDefault(StringBuffer stringBuffer, JoinPoint point) {
        //按方法限流：拼接类名 + 方法名
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());

    }


    /**
     * 按用户限流
     * @param stringBuffer
     */
    private void limitByUser(StringBuffer stringBuffer) {
            //获取当前用户
            String userSubject = UserThreadLocal.getSubject();
            User user=new User();
            try {
                user = jacksonObjectMapper.readValue(userSubject, User.class);
            } catch (JsonProcessingException e) {
                throw new RuntimeException("无法获取当前用户");
            }
            stringBuffer.append(":user:").append(user.getId());
    }


}

3. 在Redis的配置类中整合Lua语言自定义限流脚本的执行

/**
 * 限流脚本定义
 * @return
 */
@Bean
public DefaultRedisScript<Long> limitScript() {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
    redisScript.setScriptText(limitScriptText()); //加载Lua脚本
    redisScript.setResultType(Long.class); // 返回类型为Long
    return redisScript;
}

/**
 * 限流脚本
 */
private String limitScriptText() {
    return "local key = KEYS[1]\n" +
            "local count = tonumber(ARGV[1])\n" +
            "local time = tonumber(ARGV[2])\n" +
            "local current = redis.call('get', key);\n" +
            "if current and tonumber(current) > count then\n" +
            "    return tonumber(current);\n" +
            "end\n" +
            "current = redis.call('incr', key)\n" +
            "if tonumber(current) == 1 then\n" +
            "    redis.call('expire', key, time)\n" +
            "end\n" +
            "return tonumber(current);";
}

4. 限流注解的使用

案例1：登录接口IP限流

@RateLimiter( key = "login_attempt", time = 300, // 5分钟 count = 5, 
limitType = LimitType.IP ) 
@PostMapping("/login") 
public Response login(@RequestBody LoginDTO dto) { // 登录逻辑 }

案例2：API用户维度限流

@RateLimiter( key = "api_v1:data_export", time = 3600, // 1小时 count = 10,

limitType = LimitType.USER )

@GetMapping("/export")

public void exportData() { // 数据导出逻辑 }

这样就能对上述接口实现对应的限流机制了