Rate Limiting API Calls in Go

I recently needed to connect to the Desk.com API and realized that their API imposes a rate limit. There are several ways to get around this, but I decided to go with the simplest thing that would work for my code and wrote a bucket that keeps track of how many API calls I have available, and returns a “Not enough capacity” error if I can’t consume any at this time. Then in my code that makes API calls, I simply check if I can consume an API call, and if i can’t I sleep for a few seconds.

Simple enough, except keeping track of times events can be a pain. In my case I would need to keep track of how many API calls are available in my bucket, as well as when a new set was last added to the bucket. When a new API call attemps to consume an existing API call, I would need to first calculate how many many times the timed event had occured and update the remaining API calls accordingly. Only after doing these calculations could I determine if I could consume an API call.

Concurrency in Go makes this much simpler. Instead of keeping track of the last time API calls were added to my bucket, I can instead create a go routine with the sole responsibility of adding API calls to the bucket at a set interval. This ends up simplifying the code significantly, and the final rate limiter is roughly 65 lines of code. You can see it below, or you can check it out on GitHub.

package drip

import (
  "errors"
  "sync"
  "time"
)

type Bucket struct {
  Capacity     int
  DripInterval time.Duration
  PerDrip      int
  consumed     int
  started      bool
  kill         chan bool
  m            sync.Mutex
}

func (b *Bucket) Start() error {
  if b.started {
    return errors.New("Bucket was already started.")
  }

  ticker := time.NewTicker(b.DripInterval)
  b.started = true
  b.kill = make(chan bool, 1)

  go func() {
    for {
      select {
      case <-ticker.C:
        b.m.Lock()
        b.consumed -= b.PerDrip
        if b.consumed < 0 {
          b.consumed = 0
        }
        b.m.Unlock()
      case <-b.kill:
        return
      }
    }
  }()

  return nil
}

func (b *Bucket) Stop() error {
  if !b.started {
    return errors.New("Bucket was never started.")
  }

  b.kill <- true

  return nil
}

func (b *Bucket) Consume(amt int) error {
  b.m.Lock()
  defer b.m.Unlock()

  if b.Capacity-b.consumed < amt {
    return errors.New("Not enough capacity.")
  }
  b.consumed += amt
  return nil
}