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 (for example, you might want to disable it when running benchmarks or carrying out load testing, and all requests are coming 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:
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 settings 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:
package main ... func (app *application) rateLimit(next http.Handler) http.Handler { // If rate limiting is not enabled, return the next handler in the chain with // with no further action. if !app.config.limiter.enabled { return next } type client struct { limiter *rate.Limiter lastSeen time.Time } var ( mu sync.Mutex clients = make(map[string]*client) ) go func() { for { time.Sleep(time.Minute) mu.Lock() for ip, client := range clients { if time.Since(client.lastSeen) > 3*time.Minute { delete(clients, ip) } } mu.Unlock() } }() return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { ip := realip.FromRequest(r) 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"
}
}