Springboot使用AOP结合Redis+Lua脚本分布式限流

@Target(ElementType.METHOD)  
@Retention(RetentionPolicy.RUNTIME)  
@Documented  
public @interface RateLimiter {  
    long DEFAULT_REQUEST = 10;  
  
    /**  
     * max 最大请求数  
     */  
    @AliasFor("max") long value() default DEFAULT_REQUEST;  
  
    /**  
     * max 最大请求数  
     */  
    @AliasFor("value") long max() default DEFAULT_REQUEST;  
    //上方写两次是为了等值使用，添加了 AliasFor 必须通过 AnnotationUtils 获取，才会生效  
  
    /**  
     * 限流key  
     */    String key() default "";  
  
    /**  
     * 超时时长，默认1分钟  
     */  
    long timeout() default 1;  
  
    /**  
     * 超时时间单位，默认 分钟  
     */  
    TimeUnit timeUnit() default TimeUnit.MINUTES;  
}

@Slf4j  
@Aspect  
@Component  
@RequiredArgsConstructor()  
public class RateLimiterAspect {  
    private final static String SEPARATOR = ":";  
    private final static String REDIS_LIMIT_KEY_PREFIX = "limit:";  
    private final StringRedisTemplate stringRedisTemplate;  
    private final RedisScript<Long> limitRedisScript;  
  
    @Pointcut("@annotation(com.xkcoding.ratelimit.redis.annotation.RateLimiter)")  
    public void rateLimit() {  
  
    }  
  
    @Around("rateLimit()")  
    public Object pointcut(ProceedingJoinPoint point) throws Throwable {  
        MethodSignature signature = (MethodSignature) point.getSignature();  
        Method method = signature.getMethod();  
        // 通过 AnnotationUtils.findAnnotation 获取 RateLimiter 注解  
        RateLimiter rateLimiter = AnnotationUtils.findAnnotation(method, RateLimiter.class);  
        if (rateLimiter != null) {  
            String key = rateLimiter.key();  
            // 默认用类名+方法名做限流的 key 前缀  
            if (StrUtil.isBlank(key)) {  
                key = method.getDeclaringClass().getName() + StrUtil.DOT + method.getName();  
            }  
            // 最终限流的 key 为 前缀 + IP地址  
            // TODO: 此时需要考虑局域网多用户访问的情况，因此 key 后续需要加上方法参数更加合理 ，其实就是：当前限流 key 只包含了“自定义 key 或类名.方法名”和“IP 地址”，但在局域网环境下，多个用户可能共用同一个外网 IP，导致限流不够精确。为了解决这个问题，建议在 key 中加入方法参数（如用户 ID、请求参数等），这样可以更细粒度地区分不同用户或请求，避免误伤。
            key = key + SEPARATOR + IpUtil.getIpAddr();  
  
            long max = rateLimiter.max();  
            long timeout = rateLimiter.timeout();  
            TimeUnit timeUnit = rateLimiter.timeUnit();  
            boolean limited = shouldLimited(key, max, timeout, timeUnit);  
            if (limited) {  
                throw new RuntimeException("手速太快了，慢点儿吧~");  
            }  
        }  
  
        return point.proceed();  
    }  
  
    private boolean shouldLimited(String key, long max, long timeout, TimeUnit timeUnit) {  
        // 最终的 key 格式为：  
        // limit:自定义key:IP  
        // limit:类名.方法名:IP  
        key = REDIS_LIMIT_KEY_PREFIX + key;  
        // 统一使用单位毫秒  
        long ttl = timeUnit.toMillis(timeout);  
        // 当前时间毫秒数  
        long now = Instant.now().toEpochMilli();  
        long expired = now - ttl;  
        // 注意这里必须转为 String,否则会报错 java.lang.Long cannot be cast to java.lang.String  
	    Long executeTimes = stringRedisTemplate.execute(limitRedisScript, Collections.singletonList(key), now + "", ttl + "", expired + "", max + "");  
        if (executeTimes != null) {  
            if (executeTimes == 0) {  
                log.error("【{}】在单位时间 {} 毫秒内已达到访问上限，当前接口上限 {}", key, ttl, max);  
                return true;  
            } else {  
                log.info("【{}】在单位时间 {} 毫秒内访问 {} 次", key, ttl, executeTimes);  
                return false;  
            }  
        }  
        return false;  
    }  
}

// ❌ 这些操作不是原子的，存在竞态条件
stringRedisTemplate.opsForZSet().removeRangeByScore(key, 0, expired);  // 操作1
Long current = stringRedisTemplate.opsForZSet().zCard(key);             // 操作2
if (current < max) {
    stringRedisTemplate.opsForZSet().add(key, now, now);                // 操作3
    stringRedisTemplate.expire(key, ttl, TimeUnit.MILLISECONDS);        // 操作4
}

// ❌ 需要 3-4 次网络往返
stringRedisTemplate.opsForZSet().removeRangeByScore(key, 0, expired);  // 网络往返1
Long current = stringRedisTemplate.opsForZSet().zCard(key);             // 网络往返2
stringRedisTemplate.opsForZSet().add(key, now, now);                    // 网络往返3
stringRedisTemplate.expire(key, ttl, TimeUnit.MILLISECONDS);           // 网络往返4

// ❌ 复杂且性能差
String lockKey = "lock:" + key;
try {
    // 获取分布式锁
    if (redisLock.tryLock(lockKey, 100, TimeUnit.MILLISECONDS)) {
        // 执行限流逻辑
        // ...
    }
} finally {
    redisLock.unlock(lockKey);
}

// 当前逻辑：超出限制就直接抛异常
if (limited) {
    throw new RuntimeException("手速太快了，慢点儿吧~");
}

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimiter {
    // 原有参数
    long max() default 5;
    long timeout() default 1;
    TimeUnit timeUnit() default TimeUnit.MINUTES;
    
    // 新增惩罚机制参数
    boolean enablePunishment() default false;        // 是否启用惩罚
    long warningThreshold() default 10;              // 警告阈值
    long banThreshold() default 20;                  // 封禁阈值
    long banDuration() default 30;                   // 封禁时长
    TimeUnit banTimeUnit() default TimeUnit.MINUTES; // 封禁时间单位
}

@Service
public class EnhancedRateLimiterService {
    
    private static final String VIOLATION_KEY_PREFIX = "violation:";
    private static final String BAN_KEY_PREFIX = "ban:";
    
    public RateLimitResult checkRateLimit(String key, RateLimiter config) {
        // 1. 检查是否被封禁
        if (isBanned(key)) {
            return RateLimitResult.banned("用户已被封禁");
        }
        
        // 2. 执行正常限流检查
        boolean normalLimited = executeNormalRateLimit(key, config);
        
        if (!normalLimited) {
            return RateLimitResult.allowed();
        }
        
        // 3. 记录违规次数
        long violationCount = recordViolation(key);
        
        // 4. 根据违规次数分级处理
        if (config.enablePunishment()) {
            return handleViolation(key, violationCount, config);
        }
        
        return RateLimitResult.rejected("请求过于频繁");
    }
    
    private RateLimitResult handleViolation(String key, long count, RateLimiter config) {
        if (count >= config.banThreshold()) {
            // 封禁用户
            banUser(key, config.banDuration(), config.banTimeUnit());
            return RateLimitResult.banned("违规次数过多，已被封禁 " + 
                config.banDuration() + " " + config.banTimeUnit().name().toLowerCase());
                
        } else if (count >= config.warningThreshold()) {
            // 警告
            return RateLimitResult.warning("警告：继续违规将被封禁，当前违规次数：" + count);
        }
        
        return RateLimitResult.rejected("请求过于频繁，违规次数：" + count);
    }
}

@RestController
public class EnhancedTestController {
    
    @RateLimiter(
        max = 5, 
        timeout = 1, 
        timeUnit = TimeUnit.MINUTES,
        enablePunishment = true,
        warningThreshold = 3,
        banThreshold = 5,
        banDuration = 30,
        banTimeUnit = TimeUnit.MINUTES
    )
    @GetMapping("/test1")
    public ResponseEntity<Dict> test1() {
        return ResponseEntity.ok(
            Dict.create()
                .set("msg", "hello,world!")
                .set("description", "正常访问")
        );
    }
}

@Data
@AllArgsConstructor
public class RateLimitResult {
    private boolean allowed;
    private String message;
    private long remainingTime; // 剩余封禁时间（秒）
    private int violationCount; // 违规次数
    
    public static RateLimitResult allowed() {
        return new RateLimitResult(true, "允许访问", 0, 0);
    }
    
    public static RateLimitResult banned(String message) {
        return new RateLimitResult(false, message, 0, 0);
    }
    
    public static RateLimitResult warning(String message) {
        return new RateLimitResult(true, message, 0, 0);
    }
}