B站微服务框架自适应限流模块分析

// Limiter limit interface.
type Limiter interface {
  Allow(ctx context.Context) (func(Op), error)
}

// Limiter use tcp vegas + codel for adaptive limit.
type Limiter struct {
  rate  *vegas.Vegas
  queue *aqm.Queue
}

// Allow immplemnet rate.Limiter.
// if error is returned,no need to call done()
func (l *Limiter) Allow(ctx context.Context) (func(rate.Op), error) {
  var (
    done func(time.Time, rate.Op)
    err  error
    ok   bool
  )
  if done, ok = l.rate.Acquire(); !ok {
    // NOTE exceed max inflight, use queue
    if err = l.queue.Push(ctx); err != nil {
      done(time.Time{}, rate.Ignore)
      return func(rate.Op) {}, err
    }
  }
  start := time.Now()
  return func(op rate.Op) {
    done(start, op)
    l.queue.Pop()
  }, nil
}

type Config struct {
  Target   int64 // target queue delay (default 20 ms). 能容忍的最大排队延时
  Internal int64 // sliding minimum time window width (default 500 ms)  窗口大小
}

type packet struct {
  ch chan bool
  ts int64
}

// Queue queue is CoDel req buffer queue.
type Queue struct {
  pool    sync.Pool
  packets chan packet

  mux      sync.RWMutex
  conf     *Config
  count    int64
  dropping bool  //   Equal to 1 if in drop state
  faTime   int64 //  丢弃状态触发启动时间
  dropNext int64 // 下一次开始丢弃的时间
}

// Push req into CoDel request buffer queue.
// if return error is nil,the caller must call q.Done() after finish request handling
func (q *Queue) Push(ctx context.Context) (err error) {
  r := packet{
    ch: q.pool.Get().(chan bool),
    ts: time.Now().UnixNano() / int64(time.Millisecond),
  }
  select {
  case q.packets <- r:
  default:
    err = ecode.LimitExceed //channel满直接丢弃
    q.pool.Put(r.ch)
  }
  if err == nil {
    select {
    case drop := <-r.ch: //阻塞等待pop判断是否丢弃
      if drop {
        err = ecode.LimitExceed
      }
      q.pool.Put(r.ch)
    case <-ctx.Done():
      err = ecode.Deadline
    }
  }
  return
}

// Pop req from CoDel request buffer queue.
func (q *Queue) Pop() {
  for {
    select {
    case p := <-q.packets:
      drop := q.judge(p) //如果判断为丢弃，循环检测
      select {
      case p.ch <- drop:
        if !drop {
          return
        }
      default:
        q.pool.Put(p.ch)
      }
    default:
      return
    }
  }
}

func (q *Queue) controlLaw(now int64) int64 {
  q.dropNext = now + int64(float64(q.conf.Internal)/math.Sqrt(float64(q.count))) //根据当前丢弃周期内的丢包数量计算下一次丢弃的时间，丢弃的数量越多，丢弃间隔越短。
  return q.dropNext
}

// judge decide if the packet should drop or not.
func (q *Queue) judge(p packet) (drop bool) {
  now := time.Now().UnixNano() / int64(time.Millisecond)
  sojurn := now - p.ts
  q.mux.Lock()
  defer q.mux.Unlock()
  if sojurn < q.conf.Target { //延时小于容忍时间，停止丢弃
    q.faTime = 0
  } else if q.faTime == 0 { //延时较大后容忍一个窗口期
    q.faTime = now + q.conf.Internal
  } else if now >= q.faTime { //延时较大维持一个窗口期后开始进行是否进入丢弃状态的判读
    drop = true  //这里直接将返回值设为true，也就是说容忍一个窗口期后直接开始丢弃，而原始算法实现中丢弃只在状态下发生，如果最近没有发生丢弃应该继续容忍一个一个窗口期再开始丢弃的。
  }
  if q.dropping { //丢弃状态下延时较小后直接停止
    if !drop {
      // sojourn time below target - leave dropping state
      q.dropping = false
    } else if now > q.dropNext {
      q.count++
      q.dropNext = q.controlLaw(q.dropNext)
      drop = true
      return
    }   //延时较大维持一个窗口期后如果距上一次计算出的丢弃时间在一个窗口期内直接进入丢弃状态，如果一个窗口期内没有丢弃则再容忍一个窗口期
  } else if drop && (now-q.dropNext < q.conf.Internal || now-q.faTime >= q.conf.Internal) {
    q.dropping = true
    // If we're in a drop cycle, the drop rate that controlled the queue
    // on the last cycle is a good starting point to control it now.
    if now-q.dropNext < q.conf.Internal { //如果距上一次计算出的丢弃时间在一个窗口期内直接沿用之前的控制策略，否则从1开始重新调整
      if q.count > 2 {
        q.count = q.count - 2
      } else {
        q.count = 1
      }
    } else {
      q.count = 1
    }
    q.dropNext = q.controlLaw(now)
    drop = true
    return
  }
  return
}

// Vegas tcp vegas.
type Vegas struct {
  limit      int64
  inFlight   int64 //处理中的请求数
  updateTime int64
  minRTT     int64

  sample atomic.Value
  mu     sync.Mutex
  probes int64
}

//记录一个时间段内的统计信息
type sample struct {
  count       int64
  maxInFlight int64
  drop        int64
  // nanoseconds
  totalRTT int64
}

// Acquire No matter success or not,done() must be called at last.
func (v *Vegas) Acquire() (done func(time.Time, rate.Op), success bool) {
  inFlight := atomic.AddInt64(&v.inFlight, 1)
  if inFlight <= atomic.LoadInt64(&v.limit) {
    success = true
  }

  return func(start time.Time, op rate.Op) { //请求处理完后调用
    atomic.AddInt64(&v.inFlight, -1)
    if op == rate.Ignore {
      return
    }
    end := time.Now().UnixNano()
    rtt := end - start.UnixNano()

    s := v.sample.Load().(*sample)
    if op == rate.Drop {
      s.Add(rtt, inFlight, true)
    } else if op == rate.Success {
      s.Add(rtt, inFlight, false)
    }
    if end > atomic.LoadInt64(&v.updateTime) && s.Count() >= 16 {
      v.mu.Lock()
      defer v.mu.Unlock()
      if v.sample.Load().(*sample) != s { //其它协程已经重置sample，进入下一个窗口期，直接返回
        return
      }
      v.sample.Store(&sample{}) //重置sample

      lastRTT := s.RTT() //统计窗口期的平均处理时间
      if lastRTT <= 0 {
        return
      }
      updateTime := end + lastRTT*5 //设置当前窗口期大小
      if lastRTT*5 < _minWindowTime {
        updateTime = end + _minWindowTime
      } else if lastRTT*5 > _maxWindowTime {
        updateTime = end + _maxWindowTime
      }
      atomic.StoreInt64(&v.updateTime, updateTime)
      limit := atomic.LoadInt64(&v.limit)
      queue := float64(limit) * (1 - float64(v.minRTT)/float64(lastRTT)) //估算等待中的请求数量
      v.probes-- //每隔一段时间探测minRTT，如果窗口期的流量与限制差距较大则说明minRTT不准确，更新为lastRTT
      if v.probes <= 0 {
        maxFlight := s.MaxInFlight()
        if maxFlight*2 < v.limit || maxFlight <= _minLimit {
          v.probes = 3*limit + rand.Int63n(3*limit)
          v.minRTT = lastRTT
        }
      }
      if v.minRTT == 0 || lastRTT < v.minRTT {
        v.minRTT = lastRTT
      }
      var newLimit float64
      threshold := math.Sqrt(float64(limit)) / 2
      if s.Drop() {
        newLimit = float64(limit) - threshold
      } else if s.MaxInFlight()*2 < v.limit {
        return
      } else {
        if queue < threshold { //扩缩策略
          newLimit = float64(limit) + 6*threshold
        } else if queue < 2*threshold {
          newLimit = float64(limit) + 3*threshold
        } else if queue < 3*threshold {
          newLimit = float64(limit) + threshold
        } else if queue > 6*threshold {
          newLimit = float64(limit) - threshold
        } else {
          return
        }
      }
      newLimit = math.Max(_minLimit, math.Min(_maxLimit, newLimit))
      atomic.StoreInt64(&v.limit, int64(newLimit))
    }
  }, success
}