项目中我们有很多地方需要限流 比如接口限流 又分为全局接口限流 或者只针对每个用户的接口限流 可以根据不同的业务情况选择不同的方式即可
本次为了更加贴近项目实战我们采用的是redis配合lua脚本 来保证操作的原子性
而通过注解和Spring aop 能使得我们的操作更加便捷 需要限流的接口加上注解即可
首先声明一个注解
//限流注解
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter
{
/**
* 限流key
*/
public String key() default CacheConstants.RATE_LIMIT_KEY;
/**
* 限流时间,单位秒
*/
public int time() default 60;
/**
* 限流次数
*/
public int count() default 100;
/**
* 限流类型
*/
public LimitType limitType() default LimitType.DEFAULT;
}
从这个接口中可以去设置一些自定义参数 比如多少秒内可以执行多少次操作,还有具体是什么类型的限流 全局还是只针对每个用户的限流
@Aspect
@Component
public class RateLimiterAspect
{
private RedisTemplate<Object, Object> redisTemplate;
private RedisScript<Long> limitScript;
@Autowired
public void setRedisTemplate(RedisTemplate<Object, Object> redisTemplate)
{
this.redisTemplate = redisTemplate;
}
@Autowired
public void setLimitScript(RedisScript<Long> limitScript)
{
this.limitScript = limitScript;
}
@Before("@annotation(rateLimiter)")
public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable
{
int time = rateLimiter.time();
int count = rateLimiter.count();
String combineKey = getCombineKey(rateLimiter, point);
List<Object> keys = Collections.singletonList(combineKey);
try
{
Long number = redisTemplate.execute(limitScript, keys, count, time);
if (StringUtils.isNull(number) || number.intValue() > count)
{
throw new ServiceException("访问过于频繁,请稍候再试");
}
log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);
}
catch (ServiceException e)
{
throw e;
}
catch (Exception e)
{
throw new RuntimeException("服务器限流异常,请稍候再试");
}
}
public String getCombineKey(RateLimiter rateLimiter, JoinPoint point)
{
StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
if (rateLimiter.limitType() == LimitType.IP)
{
// 这里我们可以通过 实现一个Threadlocal 来获取当前登录用户的userid
stringBuffer.append(threadlocal.get()).append("-");
}
MethodSignature signature = (MethodSignature) point.getSignature();
Method method = signature.getMethod();
//获取到当前接口所在类和方法名 用来确定唯一标识
Class<?> targetClass = method.getDeclaringClass();
stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
return stringBuffer.toString();
}
}
通过String combineKey = getCombineKey(rateLimiter, point); 这个方法 你就可以拿到一个接口唯一标识 当做 key了
通过 redistemplate 的execute 方法去调用相应的 lua 脚本并传入参数
@Bean
public DefaultRedisScript<Long> limitScript()
{
DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
redisScript.setScriptText(limitScriptText());
redisScript.setResultType(Long.class);
return redisScript;
}
通过 new DefaultRedisScript 这个对象去构建lua 脚本 我们需要给它设定脚本文本 和返回值类型即可
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);";
}
}
Long number = redisTemplate.execute(limitScript, keys, count, time);
这行代码便是调用 刚才构造好的lua脚本了 keys 是一个List 集合所以可以保证key的顺序 脚本中用 KEYS[1] KEYS[2] 可以选择具体是第几个key 当然当前例子中只有一个key 所以无所谓
count 和 time 的顺序是要注意的 它的顺序就对应着脚本中 ARGV[1] ARGV[2] 的顺序 所以注意不要颠倒顺序 然后通过返回的 number去判断是否大于你设置的最大访问次数即可 大于抛出异常由前端提醒用户
这样我们就实现了通过redis 去完成限流的要求了
当然这个方式还可以去做一些其他的要求 比如秒杀这种并发量比较高的场景 ,因为redis的命令是串行化执行,再加上lua脚本的原子性 就保证了 秒杀场景下对原子性的要求 因此是非常合适去做的。
如果当前文章帮到了你,请点赞,收藏,关注 谢谢!