关于go:token令牌桶限流算法

go-tokenlimit

one of rate limit – token bucket, written in go

我的项目地址:

https://github.com/sado0823/g…

what?

之前聊bbr自适应限流的时候, 说到了 `令牌桶` 和 `漏桶` 的限流算法, 明天的配角就是 `令牌桶` - `token bucket`.

能够概括如下:
    - 令牌桶: 限流 (`Token Bucket`)
    - 漏桶: 整流 (`Leaky Bucket`)

令牌桶算法的原理是零碎会以一个恒定的速度往桶里放入令牌，而如果申请须要被解决，则须要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。如果令牌曾经满了，则抛弃令牌

why?

置信大家都很相熟, 在web服务中, 有时手动的调整流量限速的阈值,  并且容许肯定的流量峰值场景, 这种状况下`令牌桶`就能够派上用场了

限流特点是:

• 当bucket满了的时候, 将不在放入新的token, 不能超过最先限度值
• 依照肯定的速率进行token的减少
• 容许一次性耗费全副token, 即容许肯定的流量峰值存在

how?

根本思维

通过保护一个具备容量下限的token桶, 有token则申请通过, 无则失败, 一直的耗费和补充这个桶, 从而达到限流的成果

执行过程

1. 首先有一个桶, 外面装满了token, 存在一个容量下限
2. 用恒定的速率往桶中填入token, 填满则抛弃token
3. 每次申请进来, 取走一个token, 如果没有可用的token, 则拒绝请求

源码剖析

该实现为通过lua脚本, 在redis中保护桶的计数, 如果redis宕机了, 则由go自带的"golang.org/x/time/rate" 进行兜底

lua.lua

local rate = tonumber(ARGV[1])
local cap = tonumber(ARGV[2]) -- 最大容量
local now = tonumber(ARGV[3]) -- 以后工夫戳
local requested = tonumber(ARGV[4]) -- 须要去除的token数量
local fill_time = cap/rate
local ttl = math.floor(fill_time*2) -- 给肯定的过期工夫, 使得token复原平滑

-- KEYS[1]: token key 上次残余值
local last_tokens = tonumber(redis.call("get",KEYS[1]))
if last_tokens == nil then
    last_tokens = cap
end

-- KEYS[2]: token refreshed timestamp 上一次刷新token的工夫戳
local last_refreshed = tonumber(redis.call("get",KEYS[2]))
if last_refreshed == nil then
    last_refreshed = 0
end

-- 计算工夫差值
local delta = math.max(0, now-last_refreshed)
-- 工夫差值x减少速率 + 上次残余值, 计算出以后可用值
local left = math.min(cap,last_tokens+delta*rate)
local new_tokens = left
-- 可用值 > 申请值, 则通过
local allowed = left >= requested
if allowed then
  new_tokens = left - requested
end

redis.call("setex", KEYS[1], ttl, new_tokens)
redis.call("setex", KEYS[2], ttl, now)

return allowed

tokenlimit.go

主函数里应用的是一个interface, 为当前增加别的存储引擎留下了可能性 (大概率也不会用别的存储了…)

RedisI interface {
        Eval(ctx context.Context, script string, keys []string, args ...interface{}) (interface{}, error)
        IsErrNil(err error) bool
        Ping() bool
}

tokenlimit.go

最终实现限流器的构造体

Limiter struct {
        rate int // generate token number each second
        cap  int // at most token to store

        local *xrate.Limiter // limiter in process

        remote         RedisI // for distributed situation, can use redis
        tokenKey       string
        tsKey          string // timestamp key, tag get token time
        remoteMu       sync.Mutex
        remoteAlive    uint32 // ping remote server is alive or not
        monitorStarted bool
}

// New a Limiter
func New(rate, cap int, store RedisI, key string) *Limiter {
    return &Limiter{
        rate:        rate,
        cap:         cap,
        local:       xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), cap),
        remote:      store,
        tokenKey:    fmt.Sprintf(tokenFormat, key),
        tsKey:       fmt.Sprintf(timestampFormat, key),
        remoteAlive: 1,
    }
}

tokenlimit.go

新建redis存储的助手函数, 提供dsn即可

func NewStore(addr string) RedisI {
    return redisx.New(addr)
}

example

func test() {
    const (
        total = 100
        rate  = 5
        burst = 10
    )
  
  limiter := New(rate, total, NewStore("127.0.0.1:6379"), "test-token-limit")
  
  var allowed int
  for i := 0; i < total; i++ {
    
    time.Sleep(time.Second / time.Duration(total))
    
    if limiter.Allow() {
      allowed++
    }
    
  }
  
      fmt.Printf("Test_Remote: allowed:%d, burst:%d, rate:%d \n", allowed, burst, rate)
}

references

1. go-zero
2. wiki
3. token-bucket