Go 第三方库源码分析：uber-go/ratelimit

rl := ratelimit.New(100) // per second
prev := time.Now()for i := 0; i < 10; i++ {  now := rl.Take()  fmt.Println(i, now.Sub(prev))  prev = now}

// Output:// 0 0// 1 10ms// 2 10ms

example_test.golimiter_atomic.golimiter_mutexbased.goratelimit.goratelimit_bench_test.goratelimit_test.go

// New returns a Limiter that will limit to the given RPS.func New(rate int, opts ...Option) Limiter {  return newAtomicBased(rate, opts...)}

// newAtomicBased returns a new atomic based limiter.func newAtomicBased(rate int, opts ...Option) *atomicLimiter {  // TODO consider moving config building to the implementation  // independent code.  config := buildConfig(opts)  perRequest := config.per / time.Duration(rate)  l := &atomicLimiter{    perRequest: perRequest,    maxSlack:   -1 * time.Duration(config.slack) * perRequest,    clock:      config.clock,  }
  initialState := state{    last:     time.Time{},    sleepFor: 0,  }  atomic.StorePointer(&l.state, unsafe.Pointer(&initialState))  return l}

func buildConfig(opts []Option) config {  c := config{    clock: clock.New(),    slack: 10,    per:   time.Second,  }
  for _, opt := range opts {    opt.apply(&c)  }  return c}

type atomicLimiter struct {  state unsafe.Pointer  //lint:ignore U1000 Padding is unused but it is crucial to maintain performance  // of this rate limiter in case of collocation with other frequently accessed memory.  padding [56]byte // cache line size - state pointer size = 64 - 8; created to avoid false sharing.
  perRequest time.Duration  maxSlack   time.Duration  clock      Clock}

type state struct {  last     time.Time  sleepFor time.Duration}

// Limiter is used to rate-limit some process, possibly across goroutines.// The process is expected to call Take() before every iteration, which// may block to throttle the goroutine.type Limiter interface {  // Take should block to make sure that the RPS is met.  Take() time.Time}

// Take blocks to ensure that the time spent between multiple// Take calls is on average time.Second/rate.func (t *atomicLimiter) Take() time.Time {  var (    newState state    taken    bool    interval time.Duration  )  for !taken {    now := t.clock.Now()
    previousStatePointer := atomic.LoadPointer(&t.state)    oldState := (*state)(previousStatePointer)
    newState = state{      last:     now,      sleepFor: oldState.sleepFor,    }
    // If this is our first request, then we allow it.    if oldState.last.IsZero() {      taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))      continue    }
    // sleepFor calculates how much time we should sleep based on    // the perRequest budget and how long the last request took.    // Since the request may take longer than the budget, this number    // can get negative, and is summed across requests.    newState.sleepFor += t.perRequest - now.Sub(oldState.last)    // We shouldn't allow sleepFor to get too negative, since it would mean that    // a service that slowed down a lot for a short period of time would get    // a much higher RPS following that.    if newState.sleepFor < t.maxSlack {      newState.sleepFor = t.maxSlack    }    if newState.sleepFor > 0 {      newState.last = newState.last.Add(newState.sleepFor)      interval, newState.sleepFor = newState.sleepFor, 0    }    taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))  }  t.clock.Sleep(interval)  return newState.last}

type mutexLimiter struct {  sync.Mutex  last       time.Time  sleepFor   time.Duration  perRequest time.Duration  maxSlack   time.Duration  clock      Clock}

func (t *mutexLimiter) Take() time.Time {  t.Lock()  defer t.Unlock()
  now := t.clock.Now()
  // If this is our first request, then we allow it.  if t.last.IsZero() {    t.last = now    return t.last  }
  // sleepFor calculates how much time we should sleep based on  // the perRequest budget and how long the last request took.  // Since the request may take longer than the budget, this number  // can get negative, and is summed across requests.  t.sleepFor += t.perRequest - now.Sub(t.last)
  // We shouldn't allow sleepFor to get too negative, since it would mean that  // a service that slowed down a lot for a short period of time would get  // a much higher RPS following that.  if t.sleepFor < t.maxSlack {    t.sleepFor = t.maxSlack  }
  // If sleepFor is positive, then we should sleep now.  if t.sleepFor > 0 {    t.clock.Sleep(t.sleepFor)    t.last = now.Add(t.sleepFor)    t.sleepFor = 0  } else {    t.last = now  }
  return t.last}