x/rate/limit

限流是保证服务可用性和稳定的利器之一。go自带有限流器rate,我们来研读一下源码,源码较少,我们列出主要代码部分,注释说明:

// 限定速率
type Limit float64

// 常量Inf:无速率限制
const Inf = Limit(math.MaxFloat64)

// 限流器的结构体,本质就是一个令牌桶,
type Limiter struct {
  mu     sync.Mutex
  limit  Limit
  burst  int
  tokens float64
  // 记录上次 Limiter 被更新的时间
  last time.Time
  // lastEvent 记录速率受限制的时间
  lastEvent time.Time
}


// new一个新的限流器 *Limiter,r是速率,b是允许突发的令牌数量
func NewLimiter(r Limit, b int) *Limiter {
  return &Limiter{
    limit: r,
    burst: b,
  }
}

// Reservation 保存了限流器延迟后发生的事件信息
type Reservation struct {
  // 是否满足条件分配了令牌啊
  ok        bool
  // 限流器
  lim       *Limiter
  // 令牌数量
  tokens    int
  // 满足令牌发放的时间
  timeToAct time.Time
  // 令牌发放速率
  limit Limit
}


// 这是一个重要的内部函数,返回一个Reservation
func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
  // 互斥锁,保证原子性
  lim.mu.Lock()

  // 判断速率是不是最大,是的话无需限流
  if lim.limit == Inf {
    lim.mu.Unlock()
    return Reservation{
      ok:        true,
      lim:       lim,
      tokens:    n,
      timeToAct: now,
    }
  }

  // 调用 advance 方法,获取最新的时间、上一次取得令牌的时间、最新的token数量
  now, last, tokens := lim.advance(now)

  // 更新token数量
  tokens -= float64(n)

  // 计算等待时间,数量小于0表示令牌桶中木有可用的令牌了
  var waitDuration time.Duration
  if tokens < 0 {
    waitDuration = lim.limit.durationFromTokens(-tokens)
  }

  // 计算结果
  ok := n <= lim.burst && waitDuration <= maxFutureReserve

  // 初始化 Reservation
  r := Reservation{
    ok:    ok,
    lim:   lim,
    limit: lim.limit,
  }
  
  // 更新令牌桶数量和时间
  if ok {
    r.tokens = n
    r.timeToAct = now.Add(waitDuration)
  }

  // 更新限流器最新取得令牌的时间、数量、事件
  if ok {
    lim.last = now
    lim.tokens = tokens
    lim.lastEvent = r.timeToAct
  } else {
    lim.last = last
  }

  lim.mu.Unlock()
  return r
}

// 计算时间变化的 Limiter 的状态变化,得到最新的时间,最后一次取得令牌的时间和 token 令牌数量
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
  last := lim.last
  if now.Before(last) {
    last = now
  }

  maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
  elapsed := now.Sub(last)
  if elapsed > maxElapsed {
    elapsed = maxElapsed
  }

  delta := lim.limit.tokensFromDuration(elapsed)
  tokens := lim.tokens + delta
  if burst := float64(lim.burst); tokens > burst {
    tokens = burst
  }

  return now, last, tokens
}


// 这里单独放出来,是提醒一下,go也有精度缺失问题,所以复杂计算要注意
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
  sec := float64(d/time.Second) * float64(limit)
  nsec := float64(d%time.Second) * float64(limit)
  return sec + nsec/1e9
}

使用实例和demo可看:

golang 限流器


x/rate/limit
https://blog.puresai.com/2021/01/06/rate/
作者
puresai
许可协议