Building a High-Performance Token Bucket Rate Limiter in Go

As Go developers building scalable microservices, handling traffic spikes is a critical challenge. Unchecked request volumes can lead to system degradation, database connection exhaustion, or cascading failures. While there are many rate-limiting strategies, the Token Bucket Algorithm remains the gold standard for its ability to allow short bursts of traffic while enforcing long-term averages. In this post, we will implement a robust, thread-safe rate limiter in Go, optimizing for low-latency and high concurrency.

Why the Token Bucket Algorithm?

Unlike the Leaky Bucket algorithm, which enforces a strict output rate regardless of input, the Token Bucket allows a buffer of "tokens." If a client has a burst of requests, they can consume tokens up to the bucket's capacity. Once the bucket is empty, subsequent requests are rejected until tokens are refilled. This approach is ideal for APIs because it balances user experience (allowing quick actions) with system stability.

Core Implementation Strategy

To implement this in Go, we need to address three main components:

1. State Management: Tracking the number of tokens and the last refilling time.
2. Concurrency Safety: Ensuring the rate limiter works correctly in multi-threaded environments using sync.Mutex.
3. Efficiency: Avoiding unnecessary allocations and minimizing lock contention.

We will define a struct that holds the configuration (capacity and refill rate) and the current state. The key method will be Allow(), which attempts to consume a token.

Writing the Code

Here is a complete, production-ready implementation. Note the use of time.Duration for precision and the mutex to protect shared state.

package ratelimiter

import (
	"sync"
	"time"
)

// TokenBucket implements a thread-safe rate limiter.
type TokenBucket struct {
	mu         sync.Mutex
	tokens     float64
	maxTokens  float64
	refillRate float64 // tokens per second
	lastRefill time.Time
}

// NewTokenBucket creates a new rate limiter.
// capacity is the maximum number of tokens.
// refillRate is the number of tokens added per second.
func NewTokenBucket(capacity int, refillRate float64) *TokenBucket {
	return &TokenBucket{
		tokens:     float64(capacity),
		maxTokens:  float64(capacity),
		refillRate: refillRate,
		lastRefill: time.Now(),
	}
}

// Allow checks if a request is permitted and consumes a token if it is.
func (tb *TokenBucket) Allow() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	// Calculate tokens to add based on elapsed time
	now := time.Now()
	elapsed := now.Sub(tb.lastRefill).Seconds()
	tokensToAdd := elapsed * tb.refillRate

	// Update token count, capping at max capacity
	tb.tokens += tokensToAdd
	if tb.tokens > tb.maxTokens {
		tb.tokens = tb.maxTokens
	}
	tb.lastRefill = now

	// Attempt to consume a token
	if tb.tokens >= 1.0 {
		tb.tokens--
		return true
	}

	return false
}

// GetTokens returns the current number of available tokens.
func (tb *TokenBucket) GetTokens() float64 {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	return tb.tokens
}

Key Optimization Details

The implementation above uses a simple sync.Mutex. While effective, high-contention scenarios might benefit from sync.RWMutex if read operations dominate. However, since rate limiting is a write-heavy operation (checking and consuming tokens), a standard mutex is often sufficient and simpler.

Another critical detail is the calculation of tokensToAdd. By calculating tokens based on elapsed time rather than a fixed interval, we handle variable request patterns gracefully. If many tokens accumulate during idle periods, they are capped at maxTokens to prevent abuse during sudden traffic spikes.

Practical Usage in a Web Server

Integrating this into a Gin or standard HTTP server is straightforward. You can wrap your handler with the limiter:

package main

import (
	"fmt"
	"net/http"
	"your-project/ratelimiter"
)

var limiter = ratelimiter.NewTokenBucket(10, 2) // 10 burst, 2 req/sec sustained

func main() {
	http.HandleFunc("/api/data", func(w http.ResponseWriter, r *http.Request) {
		if !limiter.Allow() {
			http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
			return
		}
		fmt.Fprintln(w, "Success: Here is your data")
	})

	http.ListenAndServe(":8080", nil)
}

Conclusion

Implementing a rate limiter doesn't have to be complex. The Token Bucket algorithm, combined with Go's efficient concurrency primitives, provides a powerful tool for protecting your services. By focusing on simple data structures and minimal locking, you can achieve high throughput with low latency. This implementation serves as a solid foundation that can be extended with distributed locking (e.g., Redis) if you need to enforce rates across multiple service instances.