使用Redis的ZSet集合实现接口限流

背景

一般在设计后端接口的时候，一般都会预估接口能承受的最大流量是多少。那么如果瞬时流量超过了接口的承载力，我们就需要考虑接口做限流处理了。

限流实际上是指限制系统的输入流量和输出流量已保持系统的稳定性，防止极端条件下系统因为突然的请求激增而造成的崩溃。

思考

我们知道，Spring Cloud Alibaba有一个叫做Ribbon的组件，可以实现限流。在Ribbon中，我们可以通过设置限流策略来限制每个服务的并发访问量，防止服务被过多的请求压垮。一般限流可以通过四个思路去实现，分别是固定窗口算法、滑动窗口算法、漏桶算法、令牌桶算法。使用Redis的ZSet数据接口可以很方便的实现的接口限流。因为redis支持集群化部署，高可用易扩展，天然的支持分布式场景下各种高并发操作。下面给出一个demo实现用Redis的ZSet接口，实现华东窗口的接口限流。

滑动窗口

假如说我们想让某个接口每两秒处理5次的请求，

而如果2秒以内，第六个请求打到接口中，情况就变成了下面这种

代码讲解

demo使用Springboot框架，通过Spring AOP的自定义切面功能，实现在Mapping接口上加一个注解，实现对这个接口的限流。ZSet作为带score的有序集合，可以让接口uri作为ZSet的key（也可以做一个md5加密），value就放唯一的id，score存入当前时间戳（long类型）。代码还使用到了lua脚本，因为对于Redis而言，lua脚本的操作时原子的，这样也就避免了并发场景下的不安全问题。话不多少，先上代码。

Lua脚本：

-- 获取zset的key
local key = KEYS[1]

-- 脚本传入的限流大小
local limit = tonumber(ARGV[1])

-- 脚本传入的限流起始时间戳
local start = tonumber(ARGV[2])

-- 脚本传入的限流当前时间戳
local now = tonumber(ARGV[3])

-- 脚本传入的限流当前时间戳
local uuid = ARGV[4]

-- 获取当前流量总数
local count = tonumber(redis.call('zcount',key, start, now))

--是否超出限流值
if count + 1 >limit then
    return 1
-- 不需要限流
else
    -- 添加当前访问时间戳到zset
    redis.call('zadd', key, now, uuid)
    -- 移除时间区间以外不用的数据，不然会导致zset过大
    redis.call('zremrangeByScore',KEYS[1], 0, ARGV[2])
    return 0
end

代码接口截图：

入口控制器：

@RestController
@RequestMapping("/base")
public class LimitController {

    @RequestMapping(path = "/getUserInfo",method = RequestMethod.GET)
    @RedisLimit(path ="/base/getUserInfo" )
    @ResponseBody
    public String getUserInfo(){
        return "new UserInfo()";
    }
}

View Code

自定义注解：@RedisLimit，其中 limitvalue就是限流的请求次数，window是窗口容量，单位为秒

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RedisLimit {
    //限流请求数
    String limitvalue() default "5";
    //窗口时间
    String window() default "2";

    String path() ;
}

切面逻辑：

实现了所有带RedisLimit注解的接口都作为切入点，在before方法里去执行lua脚本，判断此次请求是否需要被限流，如果需要被限流，则抛出异常DeduplicationException，交由统一异常处理类去处理。（测试代码略显粗糙）

@Aspect
@Component
@Slf4j
public class RedisLimitAspect {

    @Autowired
    private LimitService limitService;

    @Pointcut(value = "@annotation(com.allen.deduplication.RedisLimit)")
    public void pointcut(){
    }

    @Before("pointcut()")
    public void advicBefore(JoinPoint joinPoint) throws Throwable {
        log.info("******开始进入限流判断");
        //目的：获取切入点方法上自定义RequiredLog注解中operation属性值

        //1.1获取目标对象对应的字节码对象
        Class<?> targetCls = joinPoint.getTarget().getClass();


        //1.2获取目标方法对象

        //1.2.1 获取方法签名信息从而获取方法名和参数类型
        Signature signature = joinPoint.getSignature();

        //1.2.1.1将方法签名强转成MethodSignature类型，方便调用
        MethodSignature ms = (MethodSignature) signature;

        //1.2.2通过字节码对象以及方法签名获取目标方法对象
        Method targetMethod = targetCls.getDeclaredMethod(ms.getName(), ms.getParameterTypes());

        com.allen.deduplication.RedisLimit redisLimit = targetMethod.getAnnotation(com.allen.deduplication.RedisLimit.class);
        String path = redisLimit.path();
        Integer limitvalue = Integer.valueOf(redisLimit.limitvalue());
        Integer window = Integer.valueOf(redisLimit.window());
        limitService.limit(path,limitvalue,window);
    }

View Code

Redis限流执行类：

RedisLimitServiceImpl.java

@Service
public class RedisLimitServiceImpl implements LimitService {
    @Autowired
    private StringRedisTemplate redisTemplate;

    @Autowired
    private RedisScript<Long> redisluascript;


    @Override
    public boolean limit(String key,Integer limitsize,Integer window) {
        long now = System.currentTimeMillis();
        Long needInterc = redisTemplate.execute(redisluascript, Collections.singletonList(key),String.valueOf(limitsize),String.valueOf(now-1000*window),String.valueOf(now), String.valueOf(IdUtil.getSnowflake().nextIdStr()));

        if(1==needInterc){
            throw new DeduplicationException();
        }
        return 1==needInterc;
    }
}

统一异常处理：

@RestControllerAdvice
public class ExceptionController {

    @ExceptionHandler(DeduplicationException.class)
    public String rateLimitExceptionHandler(DeduplicationException e){
        return "Your request has been restricted in flow";
    }

看一下限流的效果：

初次调用，接口正常返回值