限流的方式有很多:
1、单机模式下,可以使用AtomicInteger、RateLimiter、Semaphore。
2、分布式下,可以使用队列(如Kafka等),但是编码比较繁杂;也可以使用Nginx限流,但是属于网关层面,不能解决所有问题(如内部服务接口)。
所以,应用层也是需要做限流操作的。这里简单结合Aop+redis+lua来实现。注:如果是需要接入层先流的话,建议还是要使用nginx自带的连接数限流模块和请求限流模块。
Lua脚本:
/**
* 限流脚本
*/
private String buildLuaScript() {
return "local c" +
"\nc = redis.call('get',KEYS[1])" +
"\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
"\nreturn c;" +
"\nend" +
"\nc = redis.call('incr',KEYS[1])" +
"\nif tonumber(c) == 1 then" +
"\nredis.call('expire',KEYS[1],ARGV[2])" +
"\nend" +
"\nreturn c;";
}
1、KEYS[1]获取传入的keys参数,(这里为redis的键key)
2、ARGV[1]获取到传入的limit参数,(这里为请求的token数量,也可以理解为次数)
3、ARGV[2]获取到传入的limit参数,(这里为使用的限流key的过期时间)
限流注解,Limit
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
// 资源名称,用于描述接口功能
String name() default "";
// 资源 key
String key() default "";
// key prefix
String prefix() default "";
// 时间的,单位秒
int period();
// 限制访问次数
int count();
// 限制类型
LimitType limitType() default LimitType.CUSTOMER;
}
切面Aspect
@Aspect
@Component
public class LimitAspect {
private final RedisTemplate<Object,Object> redisTemplate;
private static final Logger logger = LoggerFactory.getLogger(LimitAspect.class);
public LimitAspect(RedisTemplate<Object,Object> redisTemplate) {
this.redisTemplate = redisTemplate;
}
@Pointcut("@annotation(com.xx.xx.Limit)")
public void pointcut() {
}
@Around("pointcut()")
public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
HttpServletRequest request = getHttpServletRequest();
MethodSignature signature = (MethodSignature) joinPoint.getSignature();
Method signatureMethod = signature.getMethod();
Limit limit = signatureMethod.getAnnotation(Limit.class);
String key = limit.key();
ImmutableList<Object> keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));
String luaScript = buildLuaScript();
RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
if (null != count && count.intValue() <= limit.count()) {
logger.info("第{}次访问key为 {},描述为 [{}] 的接口", count, keys, limit.name());
return joinPoint.proceed();
} else {
throw new BadRequestException("访问次数受限制");
}
}
public HttpServletRequest getHttpServletRequest() {
return ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
}
}
测试
/**
* @author shamee
* 限流测试
*/
@RestController
@RequestMapping("/api/limit")
public class LimitController {
private static final AtomicInteger ATOMIC_INTEGER = new AtomicInteger();
/**
* 测试限流注解,
* @param period 过期时间
* @param count 次数
* @param name 接口描述
*
* 60秒内最多只能访问 5次,保存到redis的键名为 limitP_testK,
*/
@GetMapping
@Limit(key = "testK", period = 60, count = 5, name = "test_limit", prefix = "limitP")
public int test_limit() {
return ATOMIC_INTEGER.incrementAndGet();
}
}
简单备注,以便日后查阅。