Rate limiters in Go

A rate limiter is a very important component in backend services. It can be used to limit the request rate and protect the service from being overloaded. There are many ways to implement a rate limiter, such as the sliding window method, Token Bucket, Leaky Bucket, and so on.

In fact, the Go standard library already includes an implementation of a rate limiting algorithm, namely golang.org/x/time/rate. That limiter is implemented on top of Token Bucket.

Put simply, a token bucket is a bucket of fixed size: the system puts tokens into the bucket at a constant rate, and stops for the moment once the bucket is full. Users take tokens out of the bucket, and as long as there are tokens left they can keep taking them. If no tokens are left, they have to wait until the system has placed more into the bucket.

This article mainly focuses on how to use this component in practice:

We can construct a limiter object with the following method:

limiter := NewLimiter(10, 1);

There are two parameters here:

• The first parameter is r Limit. It represents how many tokens can be produced into the token bucket per second. Limit is actually an alias for float64.
• The second parameter is b int. b represents the capacity of the token bucket. So for the example above, the limiter constructed means its token bucket size is 1, and tokens are placed into the bucket at a rate of 10 tokens per second.

Besides directly specifying the number of tokens produced per second, you can also use the Every method to specify the interval at which tokens are placed into the token bucket, for example:

limit := Every(100 * time.Millisecond);
limiter := NewLimiter(limit, 1);

The above means one token is placed into the bucket every 100ms. In essence that’s 10 produced per second.

Limiter provides three categories of methods for users to consume tokens. You can consume one token at a time, or consume several at once. Each method represents a different way of responding when there aren’t enough tokens.

Wait/WaitN

func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

Wait is actually just WaitN(ctx,1).

When using the Wait method to consume tokens, if the number of tokens in the bucket is insufficient (less than N), the Wait method will block for a while until the token requirement is satisfied. If there are enough, it returns immediately.

As you can see here, the Wait method has a context parameter. We can set the context’s Deadline or Timeout to decide the maximum duration of this Wait.

Allow/AllowN

func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool

Allow is actually just AllowN(time.Now(),1).

The AllowN method says whether, as of a certain moment, the number currently in the bucket is at least n; if so it returns true and consumes n tokens from the bucket. Otherwise it returns false without consuming any tokens.

This usually corresponds to a production scenario like this: if the request rate is too fast, some requests are simply dropped.

Reserve/ReserveN

func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation

Reserve is equivalent to ReserveN(time.Now(), 1).

ReserveN is a bit more complicated to use: once the call completes, it returns a Reservation* object regardless of whether tokens are sufficient.

You can call the object’s Delay() method, which returns the time you need to wait. If the wait time is 0, that means no waiting is needed. You must wait out that time before doing the next piece of work.

Or, if you don’t want to wait, you can call the Cancel() method, which gives the tokens back.

As a simple example, we can use the Reserve method like this.

r := lim.Reserve()
f !r.OK() {
    // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
    return
}
time.Sleep(r.Delay())
Act() // execute the related logic

Adjusting the rate dynamically

Limiter supports adjusting the rate and the bucket size:

SetLimit(Limit) changes the rate at which tokens are placed
SetBurst(int) changes the token bucket size

With these two methods, you can dynamically change the token bucket size and the rate according to your current environment and conditions, as your needs require.

Example code

package main

import (
    "context"
    "log"
    "time"

    "golang.org/x/time/rate"
)

//limit indicates how many tokens are produced per second, buret is the max number of tokens stored
//Allow checks whether a token can be obtained right now
//Wait blocks and waits until a token is obtained
//Reserve returns the wait time, then you go get the token

func main() {
    l := rate.NewLimiter(1, 5)
    log.Println(l.Limit(), l.Burst())
    for i := 0; i < 100; i++ {
        //block and wait until a token is obtained
        log.Println("before Wait")
        c, _ := context.WithTimeout(context.Background(), time.Second*2)
        if err := l.Wait(c); err != nil {
            log.Println("limiter wait err:" + err.Error())
        }
        log.Println("after Wait")

        //returns how long you need to wait before there is a new token, so you can wait that time and then run the task
        r := l.Reserve()
        log.Println("reserve Delay:", r.Delay())

        //check whether a token can be obtained right now
        a := l.Allow()
        log.Println("Allow:", a)
    }
}