基于Redis+AOP+Lua脚本实现一个服务器限流机制

后端某些接口在高并发的压力下往往会导致性能的严重下降，为了维持我们后端服务型的高性能和高可用，我们往往可以对某些接口或某些用户去设计限流机制，控制这些热点接口的访问量，我这里利用Redis的高性能优势，并整合AOP编程和引入Lua限流脚本在SpringBoot中对任意接口或某些用户实现了访问量限流的机制，其中，我这里给出了三种限流机制：用户，IP地址，全局限流

1.定义限流方式

/**
* 限流类型
* @Author GuihaoLv
*/
public enum LimitType {
/**
* 默认策略全局限流
*/
DEFAULT,

/**
* 根据请求者IP进行限流
*/
IP,

/**
* 根据请求者的用户ID进行限流
*/
USER,

/**
* 根据请求者的部门进行限流
*/
DEPT,
}

2.引入AOP，自定义限流注解和限流处理的切面类

/**
* 限流注解
* @Author GuihaoLv
*/
//实例 @RateLimiter(time = 60, count = 5, limitType = LimitType.IP) 效果：同一IP 60秒内最多允许5次登录尝试。
@Target(ElementType.METHOD) //表示该注解仅能标注在方法上，用于对具体方法进行限流控制。
@Retention(RetentionPolicy.RUNTIME) //注解在运行时保留，可通过反射机制读取注解信息，实现动态限流逻辑。
@Documented //注解信息会包含在生成的 JavaDoc 中
public @interface RateLimiter {
/**
* 限流key
*/
public String key() default RedisConstant.RATE_LIMIT_KEY;
/**
* 限流时间,单位秒
*/
public int time() default 60;

/**
* 限流次数
*/
public int count() default 100;

/**
* 限流类型
*/
public LimitType limitType() default LimitType.DEFAULT;
}

/**
* 限流处理切面
* @Author GuihaoLv
*/
@Aspect
@Component
//确保仅当配置项 spring.cache.type=redis 时，切面才会生效
@ConditionalOnProperty(prefix = "spring.cache", name = { "type" }, havingValue = "redis", matchIfMissing = false)
public class RateLimiterAspect
{
private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);

private RedisTemplate<Object, Object> redisTemplate; //redis 操作模板，用于执行 Lua 脚本和 Redis 命令

private RedisScript<Long> limitScript; //限流核心逻辑的 Lua 脚本

@Autowired
public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate)
{
this.redisTemplate = redisTemplate;
}

@Autowired
public void setLimitScript(RedisScript<Long> limitScript)
{
this.limitScript = limitScript;
}
@Autowired
private ObjectMapper jacksonObjectMapper;

@Before("@annotation(rateLimiter)")
public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable
{
//1. 获取注解参数
int time = rateLimiter.time(); //时间窗口（秒）
int count = rateLimiter.count();// 允许的请求次数

//2. 生成唯一限流 Key
String combineKey = getCombineKey(rateLimiter, point);
List<Object> keys = Collections.singletonList(combineKey);
try
{
//3. 以key为参数执行 Lua 脚本（原子性操作） keys:Redis 存储的 Key
//Lua脚本会检查Key是否存在，如果不存在则创建并设置过期时间，如果存在则递增计数器。
Long number = redisTemplate.execute(limitScript, keys, count, time);
if (StringUtils.isEmpty(number) || number.intValue() > count)
{
throw new Exception("访问过于频繁，请稍候再试");
}
log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);
} catch (RuntimeException e)
{
throw new RuntimeException("服务器限流异常，请稍候再试");
} catch (Exception e)
{
throw e;
}
}

//IP + 类名 + 方法名 的拼接方式，确保不同场景的Key不冲突。
//Key结构清晰，便于调试和监控（如通过Redis直接查看计数器）。
public String getCombineKey(RateLimiter rateLimiter, JoinPoint point)
{
StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());

switch (rateLimiter.limitType()) {
case IP:
try {
limitByIp(stringBuffer);
} catch (IOException e) {
e.printStackTrace();
} catch (InterruptedException e) {
e.printStackTrace();
}
break;
case USER:
limitByUser(stringBuffer);
break;
case DEFAULT:
limitByDefault(stringBuffer,point);
break;
}

return stringBuffer.toString();
}

/**
* 按IP限流
* @param stringBuffer
*/
private final HttpClient httpClient = HttpClient.newHttpClient();
private void limitByIp(StringBuffer stringBuffer) throws IOException, InterruptedException {
String[] services = {
"https://api.ipify.org",
"https://icanhazip.com"
};
for (String serviceUrl : services) {
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create(serviceUrl))
.GET()
.build();
HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
String ip = response.body().trim();
stringBuffer.append(":ip:").append(ip);
}
}

/**
* 全局限流
* @param stringBuffer
* @param point
*/
private void limitByDefault(StringBuffer stringBuffer, JoinPoint point) {
//按方法限流：拼接类名 + 方法名
MethodSignature signature = (MethodSignature) point.getSignature();
Method method = signature.getMethod();
Class<?> targetClass = method.getDeclaringClass();
stringBuffer.append(targetClass.getName()).append("-").append(method.getName());

}

/**
* 按用户限流
* @param stringBuffer
*/
private void limitByUser(StringBuffer stringBuffer) {
//获取当前用户
String userSubject = UserThreadLocal.getSubject();
User user=new User();
try {
user = jacksonObjectMapper.readValue(userSubject, User.class);
} catch (JsonProcessingException e) {
throw new RuntimeException("无法获取当前用户");
}
stringBuffer.append(":user:").append(user.getId());
}

}

3. 在Redis的配置类中整合Lua语言自定义限流脚本的执行

/**
* 限流脚本定义
* @return
*/
@Bean
public DefaultRedisScript<Long> limitScript() {
DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
redisScript.setScriptText(limitScriptText()); //加载Lua脚本
redisScript.setResultType(Long.class); // 返回类型为Long
return redisScript;
}

/**
* 限流脚本
*/
private String limitScriptText() {
return "local key = KEYS[1]\\n" +
"local count = tonumber(ARGV[1])\\n" +
"local time = tonumber(ARGV[2])\\n" +
"local current = redis.call('get', key);\\n" +
"if current and tonumber(current) > count then\\n" +
" return tonumber(current);\\n" +
"end\\n" +
"current = redis.call('incr', key)\\n" +
"if tonumber(current) == 1 then\\n" +
" redis.call('expire', key, time)\\n" +
"end\\n" +
"return tonumber(current);";
}

4. 限流注解的使用

案例1：登录接口IP限流

@RateLimiter( key = "login_attempt", time = 300, // 5分钟 count = 5,
limitType = LimitType.IP )
@PostMapping("/login")
public Response login(@RequestBody LoginDTO dto) { // 登录逻辑 }

案例2：API用户维度限流

@RateLimiter( key = "api_v1:data_export", time = 3600, // 1小时 count = 10,

limitType = LimitType.USER )

@GetMapping("/export")

public void exportData() { // 数据导出逻辑 }

这样就能对上述接口实现对应的限流机制了