Spring Boot 中使用 Redis + Aop 进行限流

Spring Boot 中使用 Redis 进行限流，通常你可以采用如下几种方式：

令牌桶算法（Token Bucket）
漏桶算法（Leaky Bucket）
固定窗口计数器（Fixed Window Counter）
滑动日志窗口（Sliding Log Window）

实现 Redis 限流，可以采用 Redis 提供的数据结构和功能脚本，如 Lua 脚本、Redisson 库等。以下是使用 Redis 和 Lua 脚本来实现令牌桶限流算法的示例：

步骤一：编写 Lua 脚本。

下面是一个限流的 Lua 脚本示例，实现基本的限流功能，放在Spring Boot项目下的resources目录下。

--获取KEY
local key = KEYS[1] -- 限流的 keylocal limit = tonumber(ARGV[1]) --注解标注的限流次数local curentLimit = tonumber(redis.call('get', key) or "0")if curentLimit + 1 > limit
then return 0
else-- 自增长 1redis.call('INCRBY', key, 1)-- 设置过期时间redis.call('EXPIRE', key, ARGV[2])return curentLimit + 1
end

步骤二：定义限流注解

package your.package;import java.lang.annotation.*;@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface RedisLimit {/*** 资源的key,唯一* 作用：不同的接口，不同的流量控制*/String key() default "";/*** 最多的访问限制次数*/long permitsPerSecond() default 2;/*** 过期时间也可以理解为单位时间，单位秒，默认60*/long expire() default 60;/*** 得不到令牌的提示语*/String msg() default "系统繁忙,请稍后再试.";
}

步骤三：定义Aop切面类

package your.package;import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;import javax.annotation.PostConstruct;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;/*** Limit AOP*/
@Slf4j
@Aspect
@Component
public class RedisLimitAop {@Autowiredprivate StringRedisTemplate stringRedisTemplate;@Pointcut("@annotation(your.package.RedisLimit)")private void check() {}private DefaultRedisScript<Long> redisScript;@PostConstructpublic void init() {redisScript = new DefaultRedisScript<>();redisScript.setResultType(Long.class);redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));}@Before("check()")public void before(JoinPoint joinPoint) {MethodSignature signature = (MethodSignature) joinPoint.getSignature();Method method = signature.getMethod();//拿到RedisLimit注解，如果存在则说明需要限流RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);if (redisLimit != null) {//获取redis的keyString key = redisLimit.key();String className = method.getDeclaringClass().getName();String name = method.getName();String limitKey = key + className + method.getName();log.info(limitKey);if (StringUtils.isEmpty(key)) {throw new RedisLimitException("key cannot be null");}long limit = redisLimit.permitsPerSecond();long expire = redisLimit.expire();List<String> keys = new ArrayList<>();keys.add(key);Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));log.info("Access try count is {} for key={}", count, key);if (count != null && count == 0) {log.debug("获取key失败，key为{}", key);throw new RedisLimitException(redisLimit.msg());}}}}

步骤四：自定义Redis限流异常

package your.package;/*** Redis限流自定义异常* @date 2023/3/10 21:43*/
public class RedisLimitException extends RuntimeException{public RedisLimitException(String msg) {super( msg );}
}

步骤五：自定义ResultInfo返回实体

package your.package;import lombok.Getter;
import lombok.Setter;@Getter
@Setter
public class ResultInfo<T> {private String message;private String code;private T data;public ResultInfo(String message, String code, T data) {this.message = message;this.code = code;this.data = data;}public static ResultInfo error(String message) {return new ResultInfo(message,"502",null);}}

步骤六：定义Controller接口

package your.package;import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;@Slf4j
@RestController
@RequestMapping("/limit/redis")
public class LimitRedisController {/*** 基于Redis AOP限流*/@GetMapping("/test")@RedisLimit(key = "redis-limit:test", permitsPerSecond = 2, expire = 1, msg = "当前排队人数较多，请稍后再试！")public String test() {log.info("限流成功。。。");return "ok";}}