Building a simple rate limiter in Go

Rate limiting is a critical concept in web development and API design. It ensures that users or systems can only make a limited number of requests to a server within a specific time frame. In this blog post, we’ll explore what rate limiting is, why it’s essential, and how to implement a simple rate limiter in Go.

What Is Rate Limiting?

Imagine a theme park with a roller coaster ride that can only accommodate 10 people every 10 minutes. If more than 10 people try to get on within that timeframe, they’ll have to wait. This analogy mirrors the principle of rate limiting in software systems.

In technical terms, rate limiting restricts the number of requests a client (e.g., a user, device, or IP address) can send to a server within a predefined period. It helps:

1. Prevent abuse and ensure fair usage of resources.
2. Protect servers from being overwhelmed by excessive traffic.
3. Avoid costly overuse of third-party APIs or services.

For example, an API might allow 100 requests per minute per user. If a user exceeds this limit, the server denies further requests until the limit resets.

How Does Rate Limiting Work?

One common way to implement rate limiting is through the token bucket algorithm. Here’s how it works:

1. A bucket starts with a fixed number of tokens (e.g., 10).
2. Each request removes one token from the bucket.
3. If the bucket has no tokens left, the request is denied.
4. Tokens are replenished at a steady rate (e.g., 1 token every second) until the bucket is full.

Building a Simple Rate Limiter in Go

Let’s dive into building a rate limiter in Go that limits each client to 3 requests per minute.

Step 1: Define the Rate Limiter Structure

We’ll use the sync.Mutex to ensure thread safety and store information like the number of tokens, the maximum capacity, and the refill rate.

package main

import (
    "sync"
    "time"
)

type RateLimiter struct {
    tokens         float64   // Current number of tokens
    maxTokens      float64   // Maximum tokens allowed
    refillRate     float64   // Tokens added per second
    lastRefillTime time.Time // Last time tokens were refilled
    mutex          sync.Mutex
}

func NewRateLimiter(maxTokens, refillRate float64) *RateLimiter {
    return &RateLimiter{
        tokens:         maxTokens,
        maxTokens:      maxTokens,
        refillRate:     refillRate,
        lastRefillTime: time.Now(),
    }
}

Step 2: Implement Token Refill Logic

Tokens should be replenished periodically based on the elapsed time since the last refill.

func (r *RateLimiter) refillTokens() {
    now := time.Now()
    duration := now.Sub(r.lastRefillTime).Seconds()
    tokensToAdd := duration * r.refillRate

    r.tokens += tokensToAdd
    if r.tokens > r.maxTokens {
        r.tokens = r.maxTokens
    }
    r.lastRefillTime = now
}

Step 3: Check If a Request Is Allowed

The Allow method will determine if a request can proceed based on the available tokens.

func (r *RateLimiter) Allow() bool {
    r.mutex.Lock()
    defer r.mutex.Unlock()

    r.refillTokens()

    if r.tokens >= 1 {
        r.tokens--
        return true
    }
    return false
}

Step 4: Apply Rate Limiting Per IP

To limit requests per client, we’ll create a map of IP addresses to their respective rate limiters.

type IPRateLimiter struct {
    limiters map[string]*RateLimiter
    mutex    sync.Mutex
}

func NewIPRateLimiter() *IPRateLimiter {
    return &IPRateLimiter{
        limiters: make(map[string]*RateLimiter),
    }
}

func (i *IPRateLimiter) GetLimiter(ip string) *RateLimiter {
    i.mutex.Lock()
    defer i.mutex.Unlock()

    limiter, exists := i.limiters[ip]
    if !exists {
        // Allow 3 requests per minute
        limiter = NewRateLimiter(3, 0.05)
        i.limiters[ip] = limiter
    }

    return limiter
}

Step 5: Create Middleware for Rate Limiting

Finally, we’ll create an HTTP middleware that enforces the rate limit for each client.

import (
    "fmt"
    "net"
    "net/http"
)

func RateLimitMiddleware(ipRateLimiter *IPRateLimiter, next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        ip, _, err := net.SplitHostPort(r.RemoteAddr)
        if err != nil {
            http.Error(w, "Invalid IP", http.StatusInternalServerError)
            return
        }

        limiter := ipRateLimiter.GetLimiter(ip)
        if limiter.Allow() {
            next(w, r)
        } else {
            http.Error(w, "Rate Limit Exceeded", http.StatusTooManyRequests)
        }
    }
}

Step 6: Set Up the Server

Here’s how to hook it all together and test the rate limiter.

func handleRequest(w http.ResponseWriter, _ *http.Request) {
    fmt.Fprintf(w, "Request processed successfully at %v\n", time.Now())
}

func main() {
    ipRateLimiter := NewIPRateLimiter()

    mux := http.NewServeMux()
    mux.HandleFunc("/", RateLimitMiddleware(ipRateLimiter, handleRequest))

    fmt.Println("Server starting on :8080")
    http.ListenAndServe(":8080", mux)
}

Testing the Rate Limiter

Start the server and test it using curl or your browser:

curl -X GET http://localhost:8080

• Send 3 requests quickly: All should succeed.
• Send a 4th request within the same minute: You should see Rate Limit Exceeded message.
• Wait for 20 seconds and try again: The bucket refills, and requests should succeed.

Source Code

GitHub Repo