【瞎搞瞎搞】魔改Redis实现限流

February 20, 2016

为什么有这个需求？因为我手贱。当时同事用我的手机号码测试短信下发，于是我直接用网页上的接口无限replay xhr，也骚扰一下他，这时候发现我可以无限制的发。发一条都是要钱的嘛，所以自然有了限流的需求。第一版是同事做的，直接用的Guava里面的Ratelimiter，所以并不能实现多个服务实例共享限流信息，于是需要我开发第二版限流。

限流算法就是采用的token bucket，可以应对突发情况。于是我实现了三种版本

修改Redis源码（这个时候还是Redis 3.0x）
实现一个lua脚本
Redis事务

qps自然是1 > 2 > 3 我还额外实现了和spring的集成，很轻松就可以在项目中定义ratelimiter。怀念我那些还会写java的日子。

虽然从稳定和升级的角度来看，确实不应该修改Redis源码来新增一个命令，但是我还是觉得这样魔改一下挺有意思的。

大概过程就是找到程序入口，找到get/set这种命令是如何实现的，仿照增加一个新的命令。

redis.c

struct redisCommand redisCommandTable[] = {
    {"tokenBucket",tokenBucketCommand,8,"wmF",0,NULL,1,1,1,0,0},
    {"leakBucket",leakBucketCommand,3,"wmF",0,NULL,1,1,1,0,0},
    {"get",getCommand,2,"rF",0,NULL,1,1,1,0,0},
    {"set",setCommand,-3,"wm",0,NULL,1,1,1,0,0},
    {"setnx",setnxCommand,3,"wmF",0,NULL,1,1,1,0,0},
    ...

t_string.c

void tokenBucketCommand(redisClient *c){

    long long oldTokenNum, lastCheckTime, nowTime, nowTokenNum, basicTokenNum, bucketSize, numPerPeriod, period, expiredTime;
    robj *bucketName, *bucketTime, *res, *timeObject;

    ...

    if (lastCheckTime == 0){
        nowTokenNum = basicTokenNum;
        lastCheckTime = nowTime;
    }else{
        nowTokenNum = oldTokenNum + (nowTime-lastCheckTime)/period*numPerPeriod;
        nowTokenNum = nowTokenNum < bucketSize ? nowTokenNum : bucketSize;
        if( (nowTime-lastCheckTime) >= period) lastCheckTime = nowTime;
    }

    ...

有意思的点在于描述token bucket的过程中，讲述的都是以一定的速率放入令牌，很容易的就会想到需要有一个线程来做这个事情，但是实际上这个操作可以“懒“，即在获取令牌的时候再来计算是不是真的有令牌可用就行。

看看效果

tokenbucket 命令后面跟的参数分别是bucket名字，时间存储的key，初始化令牌数，最大令牌数，添加令牌时间间隔，过期时间