2015-06-16

Rate Limiting in Go: Dissection of a Package

A few months back, I was working on a project that required some form of rate limiting to moderate sudden spikes in operations triggered by influxes of incoming events on a message bus. The solution to this problem was not to implement rate limiting but rather aggregate the events on the publisher side and emit them as batches. It got me thinking: how would I implement a rate limiter?

Well, now you can see. I uploaded my package called faregate to GitHub, with its API documented here.

I chose to implement it as a token bucket load shaper, so let's dive into by discussing its parameters:

  • The crux of this class of load shaper is single parameter R, which signifies the interval at which a single token is replenished to the bucket.
  • A secondary point of interest is L, which signifies the maximum no. of tokens the bucket may have at any period.

Using these parameters, we can model how the bucket works (ignoring the users for a moment):

There exists a bucket with at most t tokens, which may be no more than L at any given time. At a fixed period R, a single token is added to the bucket if t is less than L.

A new bucket is created by calling faregate.New. This factory applies a musing from Rob Pike about application of options. Self-standingly are the option setters faregate.TokenCount, which sets L, and faregate.TokenCount, which sets R. As a side note, I chose to have faregate.TokenCount set an initial t value—mainly because I wanted to enable the ability for the bucket's user to request more than one token at the creation of the bucket, which admittedly can lead to some small amount of initial burst. A new optionFn could be added easily that changes this behavior.

We have alluded to the bucket's users so far, so let's turn our attention to them. A user requests v tokens from the bucket. If v exceeds L, that request is said to be non-conformant and an error, since at no point may the bucket break the invariant that L ≥ t, which means t == v is unfulfillable if v > L. Everything is self-explanatory if t ≥ v, but what happens if it isn't? This turns to an explicit design decision:

The Google context package does one thing right (among many others): it employs the idioms of the language, namely Context.Done returns a <-chan struct{}, which means it may be directly select {} upon. If you've used this package, this power has spoken to you. So why the diversion? Simple: wouldn't it be great if the bucket's APIs were idiomatic to that same degree?

It turns out that can happen, and that's how I designed it. You see, Faregate.Acquire returns a channel whose closure signals v tokens have been furnished for the request. It looks something like this:

package faregate

import (
 "fmt"
 "math/rand"
 "time"
)

func Must(c <-chan struct{}, err error) <-chan struct{} {
 if err != nil {
  panic(err)
 }
 return c
}

func Example() {
 rnd := rand.New(rand.NewSource(42))

 fg, err := New(RefreshInterval(time.Second), TokenCount(100), ConcurrencyLevel(1))
 if err != nil {
  panic(err)
 }
 defer fg.Close()

 for {
  q := rnd.Intn(10)
  <-Must(fg.Acquire(uint64(q)))
  fmt.Println("acquired", q)
 }
}
(Source)

… which brings us back to that v > t case. The user can simply select {} over the result of Faregate.Acquire to determine when the tokens are available. So, speaking of context.Context, an implementation that is using them for cancellation, deadlines, etc. can use the following pattern:

package faregate

import (
 "fmt"
 "math/rand"
 "time"

 "golang.org/x/net/context"
)

func Must(c <-chan struct{}, err error) <-chan struct{} {
 if err != nil {
  panic(err)
 }
 return c
}

func Example() {
 var (
  rnd     = rand.New(rand.NewSource(42))
  bground = context.Background()
 )

 fg, err := New(RefreshInterval(time.Second), TokenCount(100), ConcurrencyLevel(1))
 if err != nil {
  panic(err)
 }
 defer fg.Close()

 for {
  var (
   q   = rnd.Intn(10)
   ctx = context.WithDeadline(bground, 5*time.Second)
  )
  select {
  case <-Must(fg.Acquire(uint64(q))):
   fmt.Println("acquired", q)
  case <-ctx.Done():
   fmt.Println("context cancelled/timedout/deadlined:", ctx.Err())
  }
 }
}
(Source)

To get an idea on why the idiomatic nature is important, replace the Faregate.Acquire API with one that blocks (e.g., Acquire(int) error as signature) or returns a truth value to indicate ability to acquire (e.g., CanAcquire(int) boolean as signature), and compare how complicated the alternatives to the original examples above become. Fluent and idiomatic API is key.

Another feature in the design that I am particularly happy about is the implementation supports fairness (in the sense of FIFO queueing requesters) if the faregate.ConcurrencyLevel has been set to a positive non-zero value, ideally to the number of no. of routines accessing the bucket at peak.

Are there any caveats with this implementation? One obvious one is if R is too frequent for the runtime's scheduler to effectively meter. I have not built a test suite to measure this, but the results would be interesting. A secondary one could be context switch and scheduling overhead, though it is my understanding that release 1.5 will mitigate much of this.

Are there any missing feature or warts? Here are the open items in my mind:

  • Formalizing the Must idiom into the public API if one is confident that he or she won't break the v > L invariant.
  • Providing an option to initialize explicitly the initial no. of tokens, because it currently defaults to L.
  • Making the API Context-aware.
  • Making the default concurrency level runtime.GOMAXPROCS.

… and if we turn our attention to the Go wiki, this mechanism is not too different from a trivial solution that it prescribes. Take your pick!

Either way, this was not meant to be a super serious exercise but rather an exploration to see what we can learn, especially in API development. Crack the source open, give it a good study, and let me know what improvements you'd make! You can find plenty of other implementations out there, so there's a lot to compare.

follow us in feedly

No comments :

Post a Comment

 

None of the content contained herein represents the views of my employer nor should it be construed in such a manner. . All content is copyright Matt T. Proud and may not be reproduced in whole or part without expressed permission.