Spring AOP切面+Redis Zset实现滑动窗口限流（Lua脚本）

限流最直接的好处防止高并发情况下因服务器资源过载导致的系统崩溃情况发生，具体来说，限流可以保证使用有限的资源提供最大化的服务能力，按照预期流量提供服务，超过的部分将会拒绝服务、降级等处理。

1、使用方式

@GetMapping("/test")
@RateLimiter(key = "123456", time = 30, count = 5, limitType = LimitType.IP, message = "测试过于频繁")
public CommonResult<String> test(){
    return CommonResult.success("操作成功");
}

2、限流类型枚举类

/**
 * 限流类型
 * @date 2023/10/17 10:00
 * @author luohao
 */
public enum LimitType {
    /**
     * 默认限流策略，针对某一个接口进行限流
     */
    DEFAULT,
    /**
     * 根据IP地址进行限流
     */
    IP;
}

3、自定义注解类

import java.lang.annotation.*;
/**
 * 滑动窗口限流注解
 * @date 2023/10/17 10:01
 * @author luohao
 */
@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
public @interface RateLimiter {
    /**
     * 定义redis中限流key前缀（默认rateLimit）
     */
    String key() default "rateLimit";

    /**
     * 限流时间，单位秒（默认60秒）
     * @return
     */
    int time() default 60;

    /**
     * 限流时间内限流次数（默认100次）
     * @return
     */
    int count() default 100;

    /**
     * 限流类型：1.限制接口访问次数 2.限制ip访问次数（默认接口访问次数）
     * @return
     */
    LimitType limitType() default LimitType.DEFAULT;

    /**
     * 错误提示信息（限流提示信息）
     * @return
     */
    String message() default ErrorMessageConstants.REDISSON_TRY_LOCK_ERROR;
}

4、代码实现切面类

import cn.hutool.extra.servlet.ServletUtil;
import com.scm.boss.common.exception.ApiException;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.concurrent.TimeUnit;

/**
 * 限流切面
 * @date 2023/10/17 10:08
 * @author luohao
 */
@Aspect
@Component
@Order(-1)
public class RateLimitAspect {

    @Value("${spring.application.name}")
    private String appName;

    @Resource
    private RedisTemplate redisTemplate;

    @Resource
    RedisScript<Long> redisScript;

    @Before("@annotation(com.scm.order.config.RateLimiter)")
    public void beforeRateLimit(JoinPoint point){
        //获取RateLimiter注解上的值
        MethodSignature methodSignature = (MethodSignature) point.getSignature();
        RateLimiter rateLimiter = AnnotationUtils.findAnnotation(methodSignature.getMethod(), RateLimiter.class);
        int time = rateLimiter.time();
        //拼接redis中的key值
        String rateKey = getRateLimitKey(rateLimiter,methodSignature);
        //限流逻辑实现
        ZSetOperations zSetOperations = redisTemplate.opsForZSet();
        //记录本次访问的时间结点
        long currentMs = System.currentTimeMillis();
        //防止一直存在于内存中
        redisTemplate.expire(rateKey, time, TimeUnit.SECONDS);
        // 移除{time}秒之前的访问记录（滑动窗口思想）
        zSetOperations.removeRangeByScore(rateKey, 0, currentMs - time * 1000);
        // 获得当前窗口内的访问记录数
        Long currCount = zSetOperations.zCard(rateKey);
        // 限流判断
        if (currCount >= rateLimiter.count()) {
            throw new ApiException(rateLimiter.message());
        }
        zSetOperations.add(rateKey, String.valueOf(currentMs), currentMs);
    }

    @Before("@annotation(com.scm.order.config.RateLimiter)")
    public void beforeRateLimitLua(JoinPoint jp){
        //获取RateLimiter注解上的值
        MethodSignature methodSignature = (MethodSignature) jp.getSignature();
        RateLimiter rateLimiter = AnnotationUtils.findAnnotation(methodSignature.getMethod(), RateLimiter.class);
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        //构建redis中的key值
        String rateKey = getRateLimitKey(rateLimiter,methodSignature);
        //记录本次访问的时间结点
        long currentMs = System.currentTimeMillis();
        //执行Lua脚本
        Long current = (Long) redisTemplate.execute(redisScript, Collections.singletonList(rateKey), count, time * 1000, currentMs);
        // 限流判断
        if(current == null || current.intValue() >= count){
            throw new ApiException(rateLimiter.message());
        }
    }

    /**
     * 拼接redis中key
     * @param rateLimiter
     * @param methodSignature
     * @return
     */
    private String getRateLimitKey(RateLimiter rateLimiter, MethodSignature methodSignature) {
        Method method = methodSignature.getMethod();
        //获取全类名
        String className = method.getDeclaringClass().getName();
        //获取方法名
        String methodName = method.getName();
        StringBuffer key = new StringBuffer(appName).append(":").append(className).append(":").append(methodName).append(":").append(rateLimiter.key());
        if(rateLimiter.limitType() == LimitType.IP){//如果参数类型为IP
            //获取客户端ip
            String clientIP = ServletUtil.getClientIP(((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest());
            key.append(":").append(clientIP);
        }
        return key.toString();
    }
}

5、保证Redis中操作原子性我们可以通过Lua脚本实现

-- 设置滑动窗口大小
local windowSize = tonumber(ARGV[1])
-- 过期时间
local expireTime = tonumber(ARGV[2])
-- 当前时间戳
local timestamp = tonumber(ARGV[3])

-- 移除过期的元素
redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, timestamp - expireTime)

-- 获取当前窗口内的元素数量
local count = redis.call('ZCARD', KEYS[1])

-- 如果当前窗口内的元素数量超过窗口大小，删除最旧的元素
if count >= windowSize then
    return count
end

-- 将新元素添加到ZSET中
redis.call('ZADD', KEYS[1], timestamp, timestamp)

--设置过期时间
redis.call('EXPIRE',KEYS[1], expireTime / 1000)

-- 返回当前窗口内的元素数量
return count

放在Redis配置类

/**
     * 定义lua脚本
     */
    @Bean
    public DefaultRedisScript<Long> limitScript(){
        DefaultRedisScript<Long> script = new DefaultRedisScript<>();
        script.setResultType(Long.class);//设置lua脚本返回值类型 需要同lua脚本中返回值一致
        script.setScriptText(limitScriptText());//读取lua脚本
        return script;
    }
    /**
     * 限流脚本
     */
    private String limitScriptText()
    {
        return "-- 设置滑动窗口大小\n" +
                "local windowSize = tonumber(ARGV[1])\n" +
                "-- 过期时间\n" +
                "local expireTime = tonumber(ARGV[2])\n" +
                "-- 当前时间戳\n" +
                "local timestamp = tonumber(ARGV[3])\n" +
                "\n" +
                "-- 移除过期的元素\n" +
                "redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, timestamp - expireTime)\n" +
                "\n" +
                "-- 获取当前窗口内的元素数量\n" +
                "local count = redis.call('ZCARD', KEYS[1])\n" +
                "\n" +
                "-- 如果当前窗口内的元素数量超过窗口大小，删除最旧的元素\n" +
                "if count >= windowSize then\n" +
                "    return count\n" +
                "end\n" +
                "\n" +
                "-- 将新元素添加到ZSET中\n" +
                "redis.call('ZADD', KEYS[1], timestamp, timestamp)\n" +
                "\n" +
                "--设置过期时间\n" +
                "redis.call('EXPIRE',KEYS[1], expireTime / 1000)\n" +
                "\n" +
                "-- 返回当前窗口内的元素数量\n" +
                "return count";
    }