Redis 限流

发布于:2025-02-22 ⋅ 阅读:(13) ⋅ 点赞:(0)
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface AccessLimit {

    /**
     * 限制次数
     */
    int count() default 15;

    /**
     * 时间窗口,单位为秒
     */
    int seconds() default 60;
}

@Aspect
@Component
public class AccessLimitAspect {

    private static final String LUA_SCRIPT =
            "local key = KEYS[1] " +
                    "local limit = tonumber(ARGV[1]) " +
                    "local current = tonumber(redis.call('get', key) or '0') " +
                    "if current + 1 > limit then " +
                    "   return 0 " +
                    "else " +
                    "   redis.call('INCR', key) " +
                    "   redis.call('EXPIRE', key, ARGV[2]) " +
                    "   return 1 " +
                    "end";

    private static final RedisScript<Long> SCRIPT_INSTANCE = new DefaultRedisScript<>(LUA_SCRIPT, Long.class);

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Before("@annotation(accessLimit)")
    public void checkAccessLimit(JoinPoint joinPoint, AccessLimit accessLimit) throws Throwable {
        validateAccessLimitParams(accessLimit);
        
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        String ipAddr = IpUtils.getIpAddr(request);
        String cacheKey = generateCacheKey(joinPoint, ipAddr);

        Long result = redisTemplate.execute(SCRIPT_INSTANCE,
                Collections.singletonList(cacheKey),
                accessLimit.count(),
                accessLimit.seconds());

        if (result != null && result == 0) {
            throw new RateLimitExceededException("操作过于频繁,请稍后再试");
        }
    }

    private void validateAccessLimitParams(AccessLimit accessLimit) {
        if (accessLimit.count() <= 0 || accessLimit.seconds() <= 0) {
            throw new IllegalArgumentException("Invalid Access Limit parameters");
        }
    }

    private String generateCacheKey(JoinPoint joinPoint, String ipAddr) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        String methodName = signature.getMethod().getName();
        return TradeCachePrefix.ACCESS_LIMIT_PREFIX + methodName + "_" + ipAddr;
    }

}

代码逻辑

  1. 定义和初始化变量
    • key = KEYS[1]:从传入的键列表中获取第一个键,作为 Redis 中存储计数器的键。
    • limit = tonumber(ARGV[1]):将传入参数中的第一个值转换为数字,这个值代表允许的最大请求次数(限制)。
    • current = tonumber(redis.call('get', key) or '0'):尝试从 Redis 中获取当前计数器的值。如果键不存在,使用默认值 '0'
  2. 判断是否超出限制
    • if current + 1 > limit then: 检查当前计数加一是否超过限制。
      • 如果超过限制,则返回 0,表示请求被拒绝。
  3. 更新计数器和设置过期时间
    • else: 如果没有超过限制:
      • redis.call('INCR', key): 将计数器加一。
      • redis.call('EXPIRE', key, ARGV[2]): 设置该键的过期时间为传入参数中的第二个值(秒)。
      • 返回 1,表示请求被接受。

应用场景

  • 限流机制:这个脚本通常用于实现基于 Redis 的限流功能。例如,在一定时间窗口内,只允许某个操作执行一定次数,以防止滥用。
  • API 请求限制:可以用于限制 API 的调用频率,每个用户或每个 IP 地址在特定时间内只能调用 API 一定次数。

使用方法

  • 执行脚本时需要传递两个参数
    • 第一个参数是允许的最大请求次数(limit)。
    • 第二个参数是键的过期时间(单位为秒),即在多长时间内重置计数器。

通过这种方式,可以有效地控制访问频率,保护系统资源不被滥用。