Implementing a Simple Rate Limiter in Go
Posted on

This is a short post about how to implement the token bucket algorithm in Go.

In a distributed system rate limiting is a core component that you should use in all services. It can help to mitigate outages when being confronted with too many requests your fleet of services could handle or making sure an internal services doesn't flood other services if they are already in a degraded state.

I will tackle two different scenarios for this example. The first one is a service simply calling another downstream service. A more complex example will take a look at the same service, that is now consuming/calling multiple downstream services.

rate-limiting

The Algorithm

We will take a look at a rather simple algorithm called Token Bucket. Imagine we are starting with a full bucket that can hold 10 tokens. Before doing a request we ask the bucket for a token, therefore each request takes 1 token out of it. After 10 requests the bucket is empty and with that we are not allowed to do any more outbound calls. Additionally we are setting up a periodic routine that fills the bucket at a constant rate, for example 1 token every second. With that we will get a new request per second once the bucket is empty. A feature of token bucket is that it can handle short bursts if the bucket is full, so a service can do more than 1 request per second for a short period.

The interface for a service implementing this can look very simple:

type Limiter interface {
    Allow(key string) bool
}

From a consumer perspective the limiter should behave as easy as possible. In the best case it could directly be added to a request library so consumers of that wouldn't even notice that they are using it.

...
if srl.Allow("service-b") {
	// do request
} else {
	// maybe log rate limited or schedule for a retry
}
...

Single Downstream Service

A single downstream service is the most straight forward case for a rate limiter. It should do exactly what I've described earlier, just make sure we don't do too many requests to the downstream service.

The implementation is also straight forward. We need to take care of two major problems:

  1. Make sure we never go over the agreed on limit. We need to protect the tokens counter against concurrent increment/decrements with a mutex.
  2. Setting up a coroutine to refill the bucket with a token at a fixed interval.

The struct for this service looks like this:

// SimpleRateLimiter implements a basic token bucket algorithm.
// It supports 'Allow' being called from multiple goroutines.
type SimpleRateLimiter struct {
	mu             sync.RWMutex
	tokens         int
	maxTokens      int
	refillDuration time.Duration

	ticker    *time.Ticker
	done      chan bool
	closeOnce sync.Once
}

The first problem is solved by the mu mutex used during the read and write operations, so when incrementing or decrementing the tokens counter.

// Allow returns true if the caller has enough tokens to be allowed to call
// another service.
// In this simple implementation the key is ignored, as there is only a single
// bucket for all potential outgoing service calls.
func (srl *SimpleRateLimiter) Allow(key string) bool {
	srl.mu.Lock()
	defer srl.mu.Unlock()

	if srl.tokens <= 0 {
		return false
	}

	srl.tokens--

	return true
}

The second problem is handled in a method called refill. We can use a ticker to get notified via an channel, when the configured duration is over and we need to act again. Additionally, we've added a done channel to exit the goroutine gracefully during shutdown.

// refill spawns a goroutine that fills the rate limiter bucket
// with one token at the configured refill duration.
func (srl *SimpleRateLimiter) refill() {
	srl.ticker = time.NewTicker(srl.refillDuration)
	srl.done = make(chan bool)

	go func() {
		for {
			select {
			case <-srl.ticker.C:
				srl.mu.Lock()
				if srl.tokens < srl.maxTokens {
					srl.tokens++
				}
				srl.mu.Unlock()

			case <-srl.done:
				srl.ticker.Stop()
				return
			}
		}
	}()
}

With that we have a simple implementation of the token bucket algorithm ready for use in a single service.

Multiple Downstream Services

The problem gets more interesting if we want to call multiple downstream services and if we want to rate limit them differently. Let's imagine a service that calls three downstream services called A, B and C. We might want different limits on service A than on service B and/or C. An easy solution could be to instantiate 3 instances of the above solution, but we want to explore the topic a bit deeper.

So to support this use case we will enhance the first solution to support multiple buckets of tokens in a single rate limiter service. To keep it simple each service will start with a full token budget. We will still stick to the same limitations of the initial solution.

Instead of keeping the tokens in a single tokens attribute of our rate limiter, we will introduce a new bucket struct that will contain them together with its own mutex. The buckets are now stored in a tokenBuckets map to support quick lookups based on the key a client provides. This key is now used to identify if we are going to downstream service A, B or C.

type bucket struct {
	tokens int
	mu     sync.RWMutex
}

// ComplexRateLimiter implements a basic token bucket algorithm.
// It supports 'Allow' being called from multiple goroutines.
type ComplexRateLimiter struct {
	mu sync.RWMutex
	// tokenBuckets keeps just a pointer to a bucket so we can work with it
	// without the need of putting it back into the map at the end of each
	// operation.
	tokenBuckets   map[string]*bucket
	maxTokens      int
	refillDuration time.Duration

	ticker    *time.Ticker
	done      chan bool
	closeOnce sync.Once
}

A key to making this performant is to lock the right structures at the right time and place. In this case the key operations are:

  1. incrementing the tokens of each bucket at the configured interval
  2. decrementing the counter of a bucket when allowing a request
  3. inserting a new bucket into our map, when we first see a client

For the first operation we will again keep it simple. All buckets will be updated in the same interval by just looping over them. This could be an issue if either the refillDuration is too frequent or there are too many buckets. We could easily enforce policies on both with checks during construction, but will just ignore them for this example. An important aspect of locks in a for loop is that you can't use defer mu.Unlock(). defer is only called when leaving a function and our for loop is no function. Therefore we must call the unlock as soon as we don't need the structure locked anymore.

// refill spawns a goroutine that fills the rate limiter bucket
// with one token at the configured refill duration.
func (srl *ComplexRateLimiter) refill() {
	srl.ticker = time.NewTicker(srl.refillDuration)
	srl.done = make(chan bool)

	go func() {
		for {
			select {
			case <-srl.ticker.C:
				// only read lock the map
				srl.mu.RLock()

				for _, bucket := range srl.tokenBuckets {
					// lock the actual bucket we want to add more tokens
					bucket.mu.Lock()
					if bucket.tokens < srl.maxTokens {
						bucket.tokens++
					}
					bucket.mu.Unlock()
				}

				srl.mu.RUnlock()

			case <-srl.done:
				srl.ticker.Stop()
				return
			}
		}
	}()
}

The Allow method gets much more interesting now. A core aspect here is actually reading the tokenBuckets twice. The first read is carried out with a read lock only and only the second one, the actual insert, will use a write lock. The reason here is performance and the expected usage pattern. I expect that we will not have hundreds of different clients, but a few static clients. And therefore I optimized for that scenario and used a read lock to get the bucket of a corresponding key out of the map fast and swallow the overhead of locking and reading again for the insert path. This snippet also handles the decrementing part which is using the mutex stored in the bucket. One thing to note here is that we don't need to store a pointer to the mutex in the bucket, because we already store the whole bucket as a reference in the map, therefore we never run into the problem of accidentially copying the mutex. Keep in mind the general rule to never copy a Mutex, only ever pass it around by pointer.

// Allow returns true if the caller has enough tokens to be allowed to call
// another service.
// Each caller is identified by the key it provides and gets its own tokens.
// If Allow is called with an unknown key a new bucket will be opened, starting
// with the configured maxTokens.
func (srl *ComplexRateLimiter) Allow(key string) bool {
	// if there is no key provided simply return false
	if len(key) == 0 {
		return false
	}

	// rlock only for the check if the key already exists in the store.
	srl.mu.RLock()
	foundBucket, found := srl.tokenBuckets[key]
	srl.mu.RUnlock()

	if !found {
		// only if we couldn't find the key do a full lock of the bucket!
		srl.mu.Lock()
		// do another lookup if the bucket wasn't inserted from a concurrent
		// routine.
		// if still not, then insert it and release the lock afterwards.
		foundBucket, found = srl.tokenBuckets[key]
		if !found {
			newBucket := &bucket{tokens: srl.maxTokens, mu: sync.RWMutex{}}
			foundBucket = newBucket
			srl.tokenBuckets[key] = newBucket
		}
		srl.mu.Unlock()
	}

	foundBucket.mu.Lock()
	defer foundBucket.mu.Unlock()

	if foundBucket.tokens <= 0 {
		return false
	}

	foundBucket.tokens--

	return true
}

Conclusion

As mentioned initially there are multiple other algorithms to solve the rate limiting problem, each of them with different ups and downs. As always the code can be found here to take a look at the other aspects of it.