SpringBoot结合Redis+Lua脚本实现api接口限流

1.场景

最近做了一个新的项目，需要提供接口供第三方调用，在api接口调用处需要实现一个限流的策略，
所以采用了 Redis + Lua脚本的一个策略来达到一个限流的目的

2.实现原理：

1. 当第三方服务在请求某一个具体的接口之前，把接口名作为key去redis中查看这个key在单位时间内的访问次数（例如1秒20次，那就设置这个key的过期时间是1秒）
2. 当这个key的次数在一秒内的次数没有达到20次，也就是没有达到限流的阈值，此时可以正常访问
3. 当这个key的次数在一秒内的次数达到了20次，也就是达到了限流的阈值，此时返回“访问频率过高,请稍后重试”的异常

3.实现步骤：

1.依赖引入：

compile ('org.springframework.boot:spring-boot-starter-data-redis')

2.redis配置类：

详细的redis配置类可以看这篇博客 Redis配置类

3.限流注解类：

/**
 * 限流参数注解
 * @author shy
 * @date 2021年1月22日 上午11:49:14
 * @param
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    /**
     * 限流唯一标示
     *
     * @return
     */
    String key() default "";

    /**
     * 限流时间
     *
     * @return
     */
    int time() default 1;

    /**
     * 限流次数
     *
     * @return
     */
    int count() default 20;
}

4.限流切面类：

结合aop，对添加限流注解的方法进行前置拦截

/**
 * 限流
 * @author shy
 * @date 2021年1月22日 下午12:16:57
 * @param
 */
@Aspect
@Configuration
public class LimitAspect {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Autowired
    private DefaultRedisScript<Number> redisluaScript;
    
    @Before("@annotation(com.huajin.cwrrapi.annotation.RateLimit)")
    public void interceptor(JoinPoint joinPoint) {
        //获取被增强的方法相关信息
		MethodSignature signature = (MethodSignature) joinPoint.getSignature();
		Method method = signature.getMethod();
		//获取方法上的RateLimit注解
		RateLimit rateLimit = method.getAnnotation(RateLimit.class);
		//获取当前请求request
		HttpServletRequest request = RequestUtil.getRequest();
		//获取第三方接口传递进来的唯一标识：appId
		String appId = request.getHeader(Constant.HEADER_APP_ID);
		if(StringUtils.isBlank(appId)) {
			 return;
		}
		StringBuilder builder = new StringBuilder(128);
		builder.append(appId).append("_").append(rateLimit.key()).append(method.getName());
		//创建单个元素的List集合 这个方法主要用于只有一个元素的优化，减少内存分配，无需分配额外的内存
		List<String> keys = Collections.singletonList(builder.toString());
		/*
		 * 通过redisTemplate来执行lua脚本
		 *      参数1：lua脚本
		 *      参数2：redis中存储的与接口名称相关的key
		 *      参数3：单位时间内的限流次数
		 *      参数4：限流的单位时间
		 */
		Number number = redisTemplate.execute(redisluaScript, keys, rateLimit.count(), rateLimit.time());
		if (number != null && number.intValue() != 0 && number.intValue() <= rateLimit.count()) {
	        return;
	    }
		throw new ErrorCodeException(ErrorCode.API_CURRENT_LIMITING);
    }

}

5.注入redisluaScript：

@Bean
    public DefaultRedisScript<Number> redisluaScript() {
        DefaultRedisScript<Number> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(“rateLimit.lua”)));
        redisScript.setResultType(Number.class);
        return redisScript;
    }

6.Lua脚本：

-- 拿到Redis中的key KEYS[1]:获取传递进来的key中的第一个key
local key = "cwrrapi_" .. KEYS[1]
-- ARGV[1]:获取传递进来的可变参数中的第一个参数 tonumber:尝试将它的参数转换为数字
-- limit:单位时间内的限制次数
local limit = tonumber(ARGV[1])
local time = tonumber(ARGV[2])
-- redis.call():在lua中执行Redis命令,
-- 获取当前key的使用次数
local current = tonumber(redis.call('get', key) or "0")
-- 如果 使用次数加一大于限制的次数,则说明达到了阈值,进行限流
if current + 1 > limit then
  return 0
else
  -- 将key中储存的数字加上指定的增量值,如果key不存在,那么key的值会先被初始化为0,然后再执行INCRBY命令
  redis.call("INCRBY", key,"1")
  -- 设置key的过期时间为1秒
  redis.call("expire", key, time)
  return current + 1
end