SpringBoot基于Redis和Lua构建分布式限流防御体系教学

-- KEYS[1..N]          N 个限流 key
-- ARGV[1..N]          N 个窗口大小（秒）
-- ARGV[N+1..2N]       N 个阈值
-- ARGV[2N+1]          当前时间戳（毫秒）
-- ARGV[2N+2]          请求唯一前缀（用于生成 member）
for i=1, #KEYS do
    local window = tonumber(ARGV[i])
    local threshold = tonumber(ARGV[i + #KEYS])
    local current = tonumber(ARGV[2 * #KEYS + 1])
    local expired = current - window * 1000
    redis.call('ZREMRANGEBYSCORE', KEYS[i], 0, expired)
    if redis.call('ZCARD', KEYS[i]) >= threshold then
        return 0
    end
end
for i=1, #KEYS do
    local window = tonumber(ARGV[i])
    local current = tonumber(ARGV[2 * #KEYS + 1])
    local member = ARGV[2 * #KEYS + 2] .. '-' .. i
    redis.call('ZADD', KEYS[i], current, member)
    redis.call('EXPIRE', KEYS[i], window)
end
return 1

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Repeatable(RateLimits.class)
public @interface RateLimit {
    /** 限流 key，支持 SpEL 表达式 */
    String key();

    /** 窗口大小（秒） */
    long window() default 60;

    /** 窗口内最大请求数 */
    long threshold() default 100;

    /** 降级方法名，必须在同类中定义，返回值与原方法一致 */
    String fallback() default "";
}

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimits {
    RateLimit[] value();
}

@Bean
public DefaultRedisScript<Long> rateLimitScript() {
    DefaultRedisScript<Long> script = new DefaultRedisScript<>();
    script.setScriptSource(new ResourceScriptSource(
        new ClassPathResource("scripts/sliding_window.lua")));
    script.setResultType(Long.class);
    return script;
}

@Aspect
@Component
public class RateLimitAspect {

    private final StringRedisTemplate redisTemplate;
    private final DefaultRedisScript<Long> rateLimitScript;
    private final SpelExpressionParser parser = new SpelExpressionParser();
    private final DefaultParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer();

    public RateLimitAspect(StringRedisTemplate redisTemplate,
                           DefaultRedisScript<Long> rateLimitScript) {
        this.redisTemplate = redisTemplate;
        this.rateLimitScript = rateLimitScript;
    }

    @Around("@annotation(io.ratelimit.RateLimit)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        RateLimit[] limits = AnnotatedElementUtils.getMergedRepeatableAnnotations(
                method, RateLimit.class, RateLimits.class);
        if (limits.length == 0) {
            return joinPoint.proceed();
        }

        List<String> keys = new ArrayList<>();
        List<Long> windows = new ArrayList<>();
        List<Long> thresholds = new ArrayList<>();

        for (RateLimit limit : limits) {
            keys.add("rate:limit:" + evalKey(limit.key(), joinPoint));
            windows.add(limit.window());
            thresholds.add(limit.threshold());
        }

        List<String> args = new ArrayList<>();
        windows.forEach(w -> args.add(String.valueOf(w)));
        thresholds.forEach(t -> args.add(String.valueOf(t)));
        args.add(String.valueOf(System.currentTimeMillis()));
        args.add(UUID.randomUUID().toString());

        Long result = redisTemplate.execute(
            rateLimitScript,
            keys,
            args.toArray(new String[0])
        );

        if (result != null && result == 1L) {
            return joinPoint.proceed();
        }

        RateLimit failed = limits[0];
        if (!failed.fallback().isEmpty()) {
            return invokeFallback(joinPoint, failed.fallback());
        }
        throw new RateLimitException("Too many requests");
    }

    private String evalKey(String expression, ProceedingJoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Object[] args = joinPoint.getArgs();
        String[] paramNames = parameterNameDiscoverer.getParameterNames(signature.getMethod());

        StandardEvaluationContext context = new StandardEvaluationContext();
        for (int i = 0; i < paramNames.length; i++) {
            context.setVariable(paramNames[i], args[i]);
        }
        return parser.parseExpression(expression).getValue(context, String.class);
    }

    private Object invokeFallback(ProceedingJoinPoint joinPoint, String fallback) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method targetMethod = joinPoint.getTarget().getClass()
                .getDeclaredMethod(fallback, signature.getParameterTypes());
        targetMethod.setAccessible(true);
        return targetMethod.invoke(joinPoint.getTarget(), joinPoint.getArgs());
    }
}

@RestController
public class SeckillController {

    @RateLimit(key = "'seckill:api'", window = 1, threshold = 1000, fallback = "apiFallback")
    @RateLimit(key = "'seckill:user:' + #userId", window = 60, threshold = 3, fallback = "userFallback")
    @RateLimit(key = "'seckill:ip:' + #ip", window = 60, threshold = 5, fallback = "ipFallback")
    @PostMapping("/seckill")
    public Order doSeckill(Long userId, String ip) {
        // 秒杀业务逻辑
        return new Order();
    }

    public Order apiFallback(Long userId, String ip) {
        return Order.failed("系统繁忙");
    }

    public Order userFallback(Long userId, String ip) {
        return Order.failed("请勿重复提交");
    }

    public Order ipFallback(Long userId, String ip) {
        return Order.failed("请求过于频繁");
    }
}

public void afterSeckillSuccess(Long userId, Long goodsId) {
    stringRedisTemplate.executePipelined((RedisCallback<Object>) connection -> {
        connection.incr("seckill:success:count".getBytes());
        connection.setEx("seckill:success:" + userId, 300, "1".getBytes());
        connection.hSet("seckill:order".getBytes(), userId.toString().getBytes(),
                goodsId.toString().getBytes());
        return null;
    });
}