Go 第三方庫(kù)源碼分析：juju/ratelimit

import "github.com/juju/ratelimit"
var tokenBucket ratelimit.Bucket = nil
func init() {  // func NewBucket(fillInterval time.Duration, capacity int64) *Bucket  // fillInterval令牌填充的時(shí)間間隔  // capacity令牌桶的最大容量  tokenBucket = ratelimit.NewBucket(200time.Millisecond, 20)}
func Handler() {  available := tokenBucket.TakeAvailable(1)  if available <= 0 {    // 限流處理  }  // handling}

ratelimit.goratelimit_test.goreader.go

func NewBucket(fillInterval time.Duration, capacity int64) *Bucket {  return NewBucketWithClock(fillInterval, capacity, nil)}

func NewBucketWithClock(fillInterval time.Duration, capacity int64, clock Clock) *Bucket {  return NewBucketWithQuantumAndClock(fillInterval, capacity, 1, clock)}

// Clock represents the passage of time in a way that// can be faked out for tests.type Clock interface {  // Now returns the current time.  Now() time.Time  // Sleep sleeps for at least the given duration.  Sleep(d time.Duration)}

// realClock implements Clock in terms of standard time functions.type realClock struct{}
// Now implements Clock.Now by calling time.Now.func (realClock) Now() time.Time {  return time.Now()}
// Now implements Clock.Sleep by calling time.Sleep.func (realClock) Sleep(d time.Duration) {  time.Sleep(d)}

func NewBucketWithQuantumAndClock(fillInterval time.Duration, capacity, quantum int64, clock Clock) *Bucket {  if clock == nil {    clock = realClock{}  }  if fillInterval <= 0 {    panic("token bucket fill interval is not > 0")  }  if capacity <= 0 {    panic("token bucket capacity is not > 0")  }  if quantum <= 0 {    panic("token bucket quantum is not > 0")  }  return &Bucket{    clock:           clock,    startTime:       clock.Now(),    latestTick:      0,    fillInterval:    fillInterval,    capacity:        capacity,    quantum:         quantum,    availableTokens: capacity,  }}

type Bucket struct {  clock Clock
  // startTime holds the moment when the bucket was  // first created and ticks began.  startTime time.Time
  // capacity holds the overall capacity of the bucket.  capacity int64
  // quantum holds how many tokens are added on  // each tick.  quantum int64
  // fillInterval holds the interval between each tick.  fillInterval time.Duration
  // mu guards the fields below it.  mu sync.Mutex
  // availableTokens holds the number of available  // tokens as of the associated latestTick.  // It will be negative when there are consumers  // waiting for tokens.  availableTokens int64
  // latestTick holds the latest tick for which  // we know the number of tokens in the bucket.  latestTick int64}

func (tb *Bucket) TakeAvailable(count int64) int64 {  tb.mu.Lock()  defer tb.mu.Unlock()  return tb.takeAvailable(tb.clock.Now(), count)}

func (tb *Bucket) takeAvailable(now time.Time, count int64) int64 {  if count <= 0 {    return 0  }  tb.adjustavailableTokens(tb.currentTick(now))  if tb.availableTokens <= 0 {    return 0  }  if count > tb.availableTokens {    count = tb.availableTokens  }  tb.availableTokens -= count  return count}

func (tb *Bucket) currentTick(now time.Time) int64 {  return int64(now.Sub(tb.startTime) / tb.fillInterval)}

func (tb *Bucket) adjustavailableTokens(tick int64) {  lastTick := tb.latestTick  tb.latestTick = tick  if tb.availableTokens >= tb.capacity {    return  }  tb.availableTokens += (tick - lastTick) * tb.quantum  if tb.availableTokens > tb.capacity {    tb.availableTokens = tb.capacity  }  return}

func (tb *Bucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool) {  tb.mu.Lock()  defer tb.mu.Unlock()  return tb.take(tb.clock.Now(), count, maxWait)}

func (tb *Bucket) take(now time.Time, count int64, maxWait time.Duration) (time.Duration, bool) {  if count <= 0 {    return 0, true  }
  tick := tb.currentTick(now)  tb.adjustavailableTokens(tick)  avail := tb.availableTokens - count  if avail >= 0 {    tb.availableTokens = avail    return 0, true  }  // Round up the missing tokens to the nearest multiple  // of quantum - the tokens won't be available until  // that tick.
  // endTick holds the tick when all the requested tokens will  // become available.  endTick := tick + (-avail+tb.quantum-1)/tb.quantum  endTime := tb.startTime.Add(time.Duration(endTick) * tb.fillInterval)  waitTime := endTime.Sub(now)  if waitTime > maxWait {    return 0, false  }  tb.availableTokens = avail  return waitTime, true}

func (tb *Bucket) Wait(count int64) {  if d := tb.Take(count); d > 0 {    tb.clock.Sleep(d)  }}

ratelimit/reader.go

func (r *reader) Read(buf []byte) (int, error) {  n, err := r.r.Read(buf)  if n <= 0 {    return n, err  }  r.bucket.Wait(int64(n))  return n, err}