Configuring the Rate Limiters

Chapter 10.3.

Configuring the Rate Limiters

At the moment our requests-per-second and burst values are hard-coded into the rateLimit() middleware. This is OK, but it would be more flexible if they were configurable at runtime instead. Likewise, it would be useful to have an easy way to turn off rate limiting altogether (which is useful when you want to run benchmarks or carry out load testing, when all requests might be coming from a small number of IP addresses).

To make these things configurable, let’s head back to our cmd/api/main.go file and update the config struct and command-line flags like so:

File: cmd/api/main.go

package main

...

type config struct {
    port int
    env  string
    db   struct {
        dsn          string
        maxOpenConns int
        maxIdleConns int
        maxIdleTime  time.Duration
    }
    // Add a new limiter struct containing fields for the requests-per-second and burst
    // values, and a boolean field which we can use to enable/disable rate limiting
    // altogether.
    limiter struct {
        rps     float64
        burst   int
        enabled bool
    }
}
...

func main() {
    var cfg config

    flag.IntVar(&cfg.port, "port", 4000, "API server port")
    flag.StringVar(&cfg.env, "env", "development", "Environment (development|staging|production)")

    flag.StringVar(&cfg.db.dsn, "db-dsn", os.Getenv("GREENLIGHT_DB_DSN"), "PostgreSQL DSN")

    flag.IntVar(&cfg.db.maxOpenConns, "db-max-open-conns", 25, "PostgreSQL max open connections")
    flag.IntVar(&cfg.db.maxIdleConns, "db-max-idle-conns", 25, "PostgreSQL max idle connections")
    flag.DurationVar(&cfg.db.maxIdleTime, "db-max-idle-time", 15*time.Minute, "PostgreSQL max connection idle time")

    // Create command line flags to read the setting values into the config struct. 
    // Notice that we use true as the default for the 'enabled' setting?
    flag.Float64Var(&cfg.limiter.rps, "limiter-rps", 2, "Rate limiter maximum requests per second")
    flag.IntVar(&cfg.limiter.burst, "limiter-burst", 4, "Rate limiter maximum burst")
    flag.BoolVar(&cfg.limiter.enabled, "limiter-enabled", true, "Enable rate limiter")

    flag.Parse()

    ...
}

...

And then let’s update our rateLimit() middleware to use these settings, like so:

File: cmd/api/middleware.go

package main

...

func (app *application) rateLimit(next http.Handler) http.Handler {
    
    ...
    
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // Only carry out the check if rate limiting is enabled.
        if app.config.limiter.enabled {
            ip, _, err := net.SplitHostPort(r.RemoteAddr)
            if err != nil {
                app.serverErrorResponse(w, r, err)
                return
            }

            mu.Lock()

            if _, found := clients[ip]; !found {
                clients[ip] = &client{
                    // Use the requests-per-second and burst values from the config
                    // struct.
                    limiter: rate.NewLimiter(rate.Limit(app.config.limiter.rps), app.config.limiter.burst),
                }
            }

            clients[ip].lastSeen = time.Now()

            if !clients[ip].limiter.Allow() {
                mu.Unlock()
                app.rateLimitExceededResponse(w, r)
                return
            }

            mu.Unlock()
        }

        next.ServeHTTP(w, r)
    })
}

Once that’s done, let’s try this out by running the API with the -limiter-burst flag and the burst value reduced to 2:

$ go run ./cmd/api/ -limiter-burst=2
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="database connection pool established"
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="starting server" addr=:4000 env=development

If you issue a batch of six requests in quick succession again, you should now find that only the first two succeed:

$ for i in {1..6}; do curl  http://localhost:4000/v1/healthcheck; done
{
    "status": "available",
    "system_info": {
        "environment": "development",
        "version": "1.0.0"
    }
}
{
    "status": "available",
    "system_info": {
        "environment": "development",
        "version": "1.0.0"
    }
}
{
    "error": "rate limit exceeded"
}
{
    "error": "rate limit exceeded"
}
{
    "error": "rate limit exceeded"
}
{
    "error": "rate limit exceeded"
}

Similarly, you can try disabling the rate limiter altogether with the -limiter-enabled=false flag like so:

$ go run ./cmd/api/ -limiter-enabled=false
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="database connection pool established"
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="starting server" addr=:4000 env=development

And you should find that all requests now complete successfully, no matter how many you make.

$ for i in {1..6}; do curl http://localhost:4000/v1/healthcheck; done
{
    "status": "available",
    "system_info": {
        "environment": "development",
        "version": "1.0.0"
    }
}
...
{
    "status": "available",
    "system_info": {
        "environment": "development",
        "version": "1.0.0"
    }
}