golang source code analysis: 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.go
limiter_atomic.go
limiter_mutexbased.go
ratelimit.go
ratelimit_bench_test.go
ratelimit_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
}