分布式令牌桶设计实现（流控）

发表于 | 分类于 ，


分布式环境下的流量控制，通过redis实现令牌桶算法。记录在设计实现过程中的思维变迁。

流程图（分布式锁实现）

流程图（分布式锁实现）

get last refill time:这里需要加入分布式锁，每次只能有1个请求获得lastRefillTime。令牌桶配置为k-v形式，lastRefillTimeLock的key和Token Bucket的key类似（比如通过前缀区分）。释放锁的时间为refill token bucket或者nowTime-lastRefillTime>refillInterval为no，更新lastRefillTime之后释放锁。

现在假设令牌桶的设置key是：tokenBucketKey，对应的分布式锁key是：lockKey。

简单说下此处分布式锁的实现：

  • set(lockKey, value, "NX", "PX", time)value为任意值，而time是毫秒数。

    每次获取lastRefillTime之前都需要通过该命令判断返回值是否为OK，OK表示获得锁。并且设置过期时间，过期表示锁自动释放。
    PS:此处还有一种基于redis的订阅/发布（sub/pub）模式的分布式锁。

  • 线程安全

    线程安全:每个JVM（进程）内的所有线程，对于获得锁应getLock到释放锁releaseLock之间的代码应加关键字synchronized，锁住这块代码避免了一个节点内的所有并发都去竞争分布式锁。（可避免分布式锁被过度竞争，将锁竞争的压力分散到各个节点中去，甚至在竞争分布式锁的时候增加竞争间隔来减轻压力）
    这有点像加个滤网，先内部竞争，优秀的人才可以去外部竞争。（后来我才知道这叫层次锁Hierarchical Lock，学至深处，万物皆然）

第一次改进（去除分布式锁）

在获取时间锁，到填充令牌的过程使用lua实现，这里利用了redis执行lua的原子性。
lua代码如下：

1
2
3
4
5
6
7
8
9
10
local nowTime = redis.call('time')[1]
local lastRefilledTime = redis.call('hget',KEYS[1],'TOKEN_BUCKET_LAST_REFILLED_TIME_OPP#10016#CMB00002')
local refilledInterval = redis.call('hget',KEYS[1],'TOKEN_BUCKET_REFILLED_INTERVAL_OPP#10016#CMB00002')
if(tonumber(nowTime) - tonumber(lastRefilledTime) > tonumber(refilledInterval))
then
  local refilledNum = redis.call('hget',KEYS[1],'TOKEN_BUCKET_REFILLED_NUM_OPP#10016#CMB00002')
  redis.call('hset',KEYS[1],'TOKEN_BUCKET_NUM_OPP#10016#CMB00002',refilledNum)
  local newRefilledTime = redis.call('time')[1]
  redis.call('hset',KEYS[1],'TOKEN_BUCKET_LAST_REFILLED_TIME_OPP#10016#CMB00002' ,newRefilledTime)
end

Redis执行lua注意要点：如上代码会报错：

1
redis.clients.jedis.exceptions.JedisDataException: ERR Error running script (call to f_f8e67ef88ffbbf7d8837a1fecb1b97f62133fea2): @user_script:7: @user_script: 7: Write commands not allowed after non deterministic commands

解析见stackoverflow
Redis tries to protect itself against such cases by blocking any write command (such as DEL) if it is executed after a random command (e.g. SCAN but also TIME, SRANDMEMBER and similar).
像TIME这种带有随机性的命令，在lua中，得到的结果不能作用于写命令。

第二次改进（将时间由外部传入）

外部传入的时间同时代表nowTime和newRefilledTime，那么newRefilledTime会比实际的newRefilledTime少。
Redis正常的情况下，误差目测为2ms以内。

第三次改进（利用redis的expire）

以上设计都不完美，再次修改lua脚本：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
private synchronized boolean refillToken(JedisCluster jedisCluster, String tokenKey, String refilledInterval, String refilledNum) {
        StringBuffer luaScript = new StringBuffer();
        luaScript.append("local nowTokenNum = redis.call('get', KEYS[1])\n")
                 .append("if not nowTokenNum then\n")
                 .append("  redis.call('set',KEYS[1],ARGV[2],'EX',ARGV[1])\n")
                 .append("  nowTokenNum=ARGV[2]\n")
                 .append("end\n")
                 .append("if(tonumber(nowTokenNum) < 0) then\n")
                 .append("   return 0\n")
                 .append("else\n")
                 .append("  nowTokenNum = redis.call('decr', KEYS[1])\n")
                 .append("  if(tonumber(nowTokenNum)<0) then\n")
                 .append("    return 0\n")
                 .append("  else\n")
                 .append("    return 1\n")
                 .append("  end\n")
                 .append("end\n");
        List<String> keys = new ArrayList<>();
        keys.add(tokenKey);
        List<String> args = new ArrayList<>();
        args.add(refilledInterval);//1
        args.add(refilledNum);//2
        Object eval = jedisCluster.eval(luaScript.toString(), keys, args);
        return eval.toString().equals("1");//eval为null抛异常
    }

很明显，如此设计去掉了时间参数的限制，老实说，我是受github上一个单机令牌桶算法的影响，跳不出【上次更新令牌时间】这个桎梏，忽略了redis的expire的特性。
大约流程为：
令牌桶（RedisExpire实现）

Redis集群的注意要点

对于一个令牌桶，涉及到的配置有：

  • 令牌桶令牌数
  • 令牌补充间隔
  • 上次补充令牌时间[利用了expire可以省掉此配置]。

如果3个配置都需要一个key，那么在redis集群对于一次补充令牌操作，便不能进行批量操作（即利用lua脚本进行令牌补充）了。
此处可设计为一个Hashmap，1个key然后3个field，便可以在集群里面将补充令牌封装成lua。
有个问题是：对于令牌数放Map里面执行decr可以使用HINCRBY -1或者封装成lua操作。
redis执行lua是原子操作，可以避免使用分布式锁。
注意：redis调用lua应该是简单的redis操作，避免阻塞其他的redis操作过久。（因为redis执行lua的原子性）

猜测一下Redis Cluster执行涉及多个key的原子操作的设计原理

  • 首先理解slot，slot是redis通过crc16(key) mod nodeNum来确定key应该在哪个节点，在redis中slot不存在物理概念，只是用来确定key所属的节点【分片】。
  • 为什么跨节点不能执行原子事务？
    猜测：redis集群是多个节点都部署了redis，节点间redis是多进程，虽然是redis是单线程，但是每个进程至少有一个线程。如果没有分布式锁，或者类似的机制，已经不能保证跨节点的原子事务。
  • 为什么必须同一个slot的key才允许原子操作？
    猜测：同一个节点的确可以保证原子性，但是如果2个key落在同一个节点的不同的slot，如果以后扩充或者减少redis集群的节点，都会导致节点拥有的slot的改变，从而导致原来同一个节点的key分片到不同节点。这样一来，坑就大了。集群节点数量改变会导致原来的代码挂掉。

写在最后

强烈建议使用Redis单实例来做，依赖可以避免Redis的事务问题；二来Redis单实例效率比集群高不少（毕竟少了一次Hash定位）。