redis实现的简单令牌桶

Posted 2021-11-27 albizzia

这里给出的令牌桶是以redis单节点为中间件, 改成以redis集群为中间件应该也很简单. 不过, 这里的实现比较简单, 主要提供两个函数, 一个用于消费令牌, 一个用于添加令牌. 这里, 消费令牌和添加令牌都是通过lua来保证原子性.

消费令牌的代码如下 :

// FetchToken 用来获取某个key的一个令牌
func (acc *Accessor) FetchToken(key string) (bool, error) {
    /*
     * KEYS[1] 表示特定的key, 这个key是当前的令牌数
     */
    keyFetchScript :=
        `--[[测试显示, 通过call, 可以将error返回给客户端, 即使没有使用return]]--
        local curNum = redis.call("DECR", KEYS[1])
        if (curNum >= 0)
        then
            return true
        end

        redis.call("INCR", KEYS[1])
        return false
        `

    keyFetchCmd := redis.NewScript(keyFetchScript)
    res, err := keyFetchCmd.Run(acc.client, []string{key}).Result()
    if err != nil && err != redis.Nil {
        return false, err
    }

    if err == redis.Nil {
        return false, nil
    }

    if val, ok := res.(int64); ok {
        return (val == 1), nil
    }

    return false, errors.New("res should be bool")
}

 这里每一个key都有一个辅助的key_idx, 每次增加key的令牌数, 都会使key_idx的值加1, 同时这个函数调用会返回对应的key_idx的值. 如果传入的idx的值与key_idx值不相同, 则不会执行增加令牌数的操作. 这样设计的目的是, 如果你在不同机器中启动多个增加令牌数的程序, 而且这些程序启动时间不同, 那么其中一个程序将会起到增加令牌数的效果, 而另外的程序不会新增令牌数. 当增加令牌数的这个程序意外关闭, 将会有新的增加令牌的程序起作用. 这个实现的思想类似于乐观锁.具体代码如下:

// AddToken 用来添加某个key的令牌
func (acc *Accessor) AddToken(key string, keyIdx int, keyAdd int, keyLimit int) (int, error) {
    /* KEYS[1] 表示特定key,这个key是当前的令牌
     * KEYS[2] 表示特定key的idx
     * ARGV[1] 表示修改的key的增加的值
     * ARGV[2] 表示修改的key的最大值
     * ARGV[3] 表示修改的key的idx的序号
     */
    // 实现思路, 先判断这个key当前的序号与修改调用的序号是否一致,如果一致, 则进行修改,否则返回当前的序号
    keyAddScript :=
        `--[[测试显示, 通过call, 可以将error返回给客户端, 即使没有使用return]]--
        local curIdx = redis.call("INCR", KEYS[2])
        if (curIdx ~= (ARGV[3]+1))
        then
            curIdx = redis.call("DECR", KEYS[2])
            return curIdx
        end
        local curNum = redis.call("INCRBY", KEYS[1], ARGV[1])
        local maxNum = tonumber(ARGV[2])
        if (curNum > maxNum) 
        then 
            redis.call("SET", KEYS[1], ARGV[2])
        end
        return curIdx
        `
    keyAddCmd := redis.NewScript(keyAddScript)
    res, err := keyAddCmd.Run(acc.client, []string{key, getKeyIdx(key)},
        keyAdd, keyLimit, keyIdx).Result()
    if err != nil && err != redis.Nil {
        return 0, err
    }

    if idx, ok := res.(int64); ok {
        return int(idx), nil
    }

    return 0, errors.New("res should be integer")
}

 完整的代码请参考如下地址:

 https://github.com/ss-torres/ratelimiter.git

在这个git地址中, 如果想调用db_test.go中的测试, 可以参考如下命令:

go test ratelimiter/db -args "localhost:6379" "" "hello"

如果有什么好的建议, 或者有什么问题, 欢迎提出