自定义注解+Lua脚本实现限流

1. 固定窗口限流（推荐首选）

适用场景：接口每秒 / 每分钟最多允许 N 次请求

优点：简单、性能极高、内存占用小

缺点：窗口边界可能有瞬间突刺（大多数业务可忽略）

Lua脚本

java 复制代码

        // 固定窗口限流
        String luaScript =
                "local key = KEYS[1]\n" +        // 限流key
                        "local limit = tonumber(ARGV[2])\n" +  // 最大请求数
                        "local window = tonumber(ARGV[3])\n" + // 窗口大小（秒）
                        "local count = redis.call('get', key)\n" + // 获取当前计数
                        "if count and tonumber(count) >= limit then\n" +
                        "    return 0\n" +              // 超过限流，返回0
                        "end\n" +
                        "count = redis.call('incr', key)\n" +  // 计数+1
                        "if count == 1 then\n" +
                        "    redis.call('expire', key, window)\n" + // 第一次设置过期时间
                        "end\n" +
                        "return 1\n";  // 允许通过，返回1

2. 滑动窗口限流（更精准）

适用场景 ：金融、支付等绝对不能超限流的场景

优点：无窗口边界突刺，限流最精准

缺点：占用稍多内存，复杂度略高

Lua脚本

java 复制代码

        // 滑动窗口限流
        String luaScript =
                "local key = KEYS[1]\n" +
                        "local now = tonumber(ARGV[1])\n" +
                        "local limit = tonumber(ARGV[2])\n" +
                        "local window = tonumber(ARGV[3])\n" +
                        "local window_start = now - window * 1000\n" +
                        "redis.call('ZREMRANGEBYSCORE', key, '-inf', window_start)\n" +  // 删除窗口外的旧请求
                        "local current = redis.call('ZCARD', key)\n" +  // 统计当前窗口内有多少请求
                        "if current < limit then\n" +    // 判断是否超过限流
                        "    local unique_id = ARGV[4]\n" +
                        "    redis.call('ZADD', key, now, unique_id)\n" +
                        "    redis.call('EXPIRE', key, window + 10)\n" +  //-- 自动过期，防残留
                        "    return 1\n" +  // 允许通过
                        "else\n" +
                        "    return 0\n" +  // 被限流
                        "end";

类型	Redis 存什么	数据结构	内存占用	性能	精准度
固定窗口	一个数字（计数器）	String	极小	极快	够用（99% 场景）
滑动窗口	一堆时间戳	ZSet	较大	较快	极高（金融 / 支付）

自定义注解

java 复制代码

/**
 * 限流注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    String limitConfig(); // 限制次数
    String windowConfig(); // 时间窗口(单位秒)
    String key() default ""; // 限制key
}

实现限流逻辑

java 复制代码

import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.core.env.Environment;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;
import java.util.UUID;

/**
 * 限流切面
 */
@Aspect
@Component
@Slf4j
public class RateLimitAspect {

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    @Resource
    private Environment environment;

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        String key = rateLimit.key();
        if (key.isEmpty()) {
            key = joinPoint.getSignature().getName();
            log.info("限流key为空, 使用方法名作为限流key: {}", key);
        }

        // 从注解参数动态读取配置值（实际项目中应该从配置中心读取）
        int limit = environment.getProperty(rateLimit.limitConfig(), Integer.class, 3);
        int window = environment.getProperty(rateLimit.windowConfig(), Integer.class, 60);
        //int limit = ConfigService.getConfig("demo-gateway.properties")
                //.getIntProperty(rateLimit.limitConfig(), 100);
        //int window = ConfigService.getConfig("demo-gateway.properties")
                //.getIntProperty(rateLimit.windowConfig(), 60);
        
        // 固定/滑动窗口限流
        String luaScript = "lua脚本"; 

        // 生成唯一标识符
        String uniqueId = UUID.randomUUID().toString();
        long currentTime = System.currentTimeMillis();

        // 执行Lua脚本
        Long result = stringRedisTemplate.execute(
                new DefaultRedisScript<>(luaScript, Long.class),
                Collections.singletonList(key),  // KEYS[1] - 限流的key
                String.valueOf(currentTime),     // ARGV[1] - 当前时间戳（毫秒）
                String.valueOf(limit), // ARGV[2] - 最大允许的请求数
                String.valueOf(window), // ARGV[3] - 时间窗口大小（秒）
                uniqueId                         // ARGV[4] - 唯一标识
        );

        if (result == 0) {
            log.warn("接口: {} 被限流", joinPoint.getSignature().getName());
            Object[] args = joinPoint.getArgs();
            for (Object arg : args) {
                // 根据实际业务需要，触发限流修改入参
                if (arg instanceof KnowledgeBody) {
                    ((KnowledgeBody) arg).setRerank(true);
                }
            }
            // 修改参数后继续执行原方法（Controller方法）
            return joinPoint.proceed(args);
        }

        return joinPoint.proceed();
    }
}