Redis+Lua脚本实现分布式服务的限流

背景

​ 限流的目的是通过对并发访问/请求进行限速或者一个时间窗口内的的请求进行限速来保护系统，一旦达到限制速率则可以拒绝服务。

​ 开始打算使用Guava RateLimiter来实现限流，但RateLimiter是局限于单机中使用，然后打算使用Redis+Lua脚本实现限流。

1 提供调用的接口

@Slf4j
@RestController
@RequestMapping("/rateLimter")
public class RateLimterController {

    @PostMapping(value = "rateLimter")
    @RateLimiter(key = "rateLimter", limit = "30", expire = "1", message = "请稍后再试")
    public void rateLimter() {
       log.info("通过测试")
    }
}

2 自定义注解

/**
 * @className RateLimiter
 * @desc 限流注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter {

    /**
     * 限流key
     * @return
     */
    String key() default "rate:limiter";
    /**
     * 单位时间限制通过请求数
     * @return
     */
    String limit() default "30";

    /**
     * 过期时间，单位秒
     * @return
     */
    String expire() default "1";

    /**
     * 限流提示语
     * @return
     */
    String message() default "请稍后再试";
}

3 设置切面

@Slf4j
@Aspect
@Component
@EnableAspectJAutoProxy
public class RateLimterAspect {


    @Autowired
    RedisCache redisTemplate;

    private DefaultRedisScript<Long> redisScript;
    private static AtomicReference<String> LUA_SHA = new AtomicReference<>();

    @PostConstruct
    public void init() {
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimter.lua")));
        log.info("RateLimterHandler[分布式限流处理器]脚本加载完成");
        //缓存lua脚本
        LUA_SHA.compareAndSet(null, redisTemplate.getRedisScriptingCommands().scriptLoad(redisScript.getScriptAsString()));
    }

    @Pointcut("@annotation(com.annotation.RateLimiter)")
    public void rateLimiter() {
    }

    @Around("@annotation(rateLimiter)")
    public Object around(ProceedingJoinPoint proceedingJoinPoint, RateLimiter rateLimiter) throws Throwable {
        if (log.isDebugEnabled()) {
            log.debug("RateLimterHandler[分布式限流处理器]开始执行限流操作");
        }
        Signature signature = proceedingJoinPoint.getSignature();
        if (!(signature instanceof MethodSignature)) {
            throw new IllegalArgumentException("the Annotation @RateLimter must used on method!");
        }
        // 限流模块key
        String limitKey = rateLimiter.key();
        Preconditions.checkNotNull(limitKey);
        // 限流阈值
        String limitTimes = rateLimiter.limit();
        // 限流超时时间
        String expireTime = rateLimiter.expire();

        if (log.isDebugEnabled()) {
            log.debug("RateLimterHandler[分布式限流处理器]参数值为-limitTimes={},expire={}", limitTimes, expireTime);
        }
        // 限流提示语
        String message = rateLimiter.message();
        //执行Lua脚本
        List<String> keyList = new ArrayList();
        // 设置key值为注解中的值
        keyList.add(limitKey);
        //调用脚本并执行
        Long result = redisTemplate.getRedisScriptingCommands().evalsha(LUA_SHA.get(),
                ScriptOutputType.INTEGER,
                keyList.toArray(new String[keyList.size()]),
                limitTimes,expireTime);
        if (Objects.equals(result,0L)) {
            log.info("由于在{}秒内，超过了最大允许的请求次数{},触发限流",expireTime,limitTimes);
            return new Response(ResponseCode.CORRECT_EXPERIENCE_FAIL,message);
        }
        if (log.isDebugEnabled()) {
            log.debug("RateLimterHandler[分布式限流处理器]限流执行结果-result={},请求[正常]响应", result);
        }
        return proceedingJoinPoint.proceed();
    }

}

4 Lua脚本

rateLimter.lua脚本如下：

--获取KEY
local key1 = KEYS[1]

local val = redis.call('incr', KEYS[1])
local ttl = redis.call('ttl', KEYS[1])

--获取ARGV内的参数并打印
local times = ARGV[1]
local expire = ARGV[2]

redis.log(redis.LOG_DEBUG,tostring(times))
redis.log(redis.LOG_DEBUG,tostring(expire))

redis.log(redis.LOG_NOTICE, "incr "..key1.." "..val);
if val == 1 then
    redis.call('expire', KEYS[1], tonumber(expire))
else
    if ttl == -1 then
        redis.call('expire', KEYS[1], tonumber(expire))
    end
end

if val > tonumber(times) then
    return 0
end

return 1

注意

在使用Lua脚本的是时候，如果Redis使用的是阿里云的集群，这里有个巨坑，大家请勿踩雷，可以参考我的另外一篇文章。
线上踩坑：Redis集群调用Lua脚本-ERR bad lua script for redis cluster, all the keys that the script uses should