View : 113

28/07/2026 14:14pm

Token Bucket Algorithm Diagram for Rate Limiting in Go

Golang The Series EP.164: Rate Limiting AI Requests - Preventing System Crashes from API Overuse

#Rate Limiting

#Token Bucket

#Golang Middleware

#Go

#Golang

#AI API

#AI Backend

Hello Gophers! In the last episode, we set up a caching system to significantly cut down repetitive workloads and slash token costs. However, when scaling a system architecture for a high volume of users, you will inevitably encounter abuse and overuse scenarios.

Imagine a user writing a script that loops queries to your AI model, or a rogue internal service bug triggering hundreds of API calls per second. External providers like OpenAI will instantly hit you with a temporary ban for exceeding RPM (Requests Per Minute) limits, bringing down the service for your entire organization. Alternatively, if you host open-source models locally, this level of spam can instantly exhaust your GPU's VRAM (Out-of-Memory), crashing your infrastructure completely.

To protect our infrastructure and ensure fair resource distribution, we must place a Rate Limiter Middleware right at the entry point of our Go server.

How the Token Bucket Algorithm Works

One of the most popular and efficient algorithms for rate limiting in Go is the Token Bucket algorithm. The core concept relies on a bucket that holds a fixed maximum capacity of tokens (e.g., a maximum of 3 tokens). A background job continuously replenishes tokens into the bucket at a stable rate (e.g., 1 token per second).

When an incoming request hits the server, the system checks if there is a token available in the bucket. If a token is present, it consumes 1 token and forwards the request to the AI Handler. If the bucket is completely empty, the system drops the request and immediately fires back an HTTP status code 429 Too Many Requests. The main advantage of this approach is its ability to handle sudden burst traffic up to the maximum capacity configuration of your bucket.

Plaintext

[ Constant Token Refill Rate ] -> (e.g., 1 token / sec)
                                   │
                                   ▼
                        ┌──────────────────┐
                        │  Bucket (Max=3)  │  <- Accumulated tokens
                        └──────────────────┘
                                   │
       [ Incoming Request ] -------┘
               │
               ▼
     (Tokens available?)
      ├───> Yes ───> Consume 1 token ───> Forward to AI Handler
      └───> No  ───> Drop Request ─────> Return HTTP 429 Too Many Requests

Installing Dependencies

We will utilize Go's official sub-package x/time/rate. This package is natively engineered to be Thread-safe, handling synchronization for high-concurrency environments flawlessly out of the box:

Bash

go get golang.org/x/time/rate
go get github.com/gin-gonic/gin

Data Structure and Client-Specific Rate Limiting

To keep the project clean, we will decouple our data states into modules. We will also implement a sync.RWMutex to completely eliminate any Data Race vulnerabilities between our main middleware logic and our background memory cleanup worker.

Part 1: State Management & Background Cleanup

Go

package main

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

// clientLimiter holds the unique token bucket state for each separate client
type clientLimiter struct {
	limiter  *rate.Limiter
	mu       sync.RWMutex // Mutex to prevent Data Races when reading/writing lastSeen
	lastSeen time.Time
}

var (
	// Use sync.Map to safely support concurrent operations across multiple Goroutines
	limiters sync.Map 
)

// init spawns a background worker to continuously scan and evict stale keys, preventing memory leaks
func init() {
	go func() {
		for {
			time.Sleep(10 * time.Minute) // Scan execution every 10 minutes
			limiters.Range(func(key, value any) bool {
				client := value.(*clientLimiter)
				
				client.mu.RLock()
				lastSeen := client.lastSeen
				client.mu.RUnlock()

				// If a client has been inactive for over 1 hour, purge them from memory
				if time.Since(lastSeen) > 1*time.Hour {
					limiters.Delete(key)
				}
				return true
			})
		}
	}()
}

Part 2: Middleware Integration & Resolver Logic

Go

// getLimiter resolves or creates a unique Rate Limiter instance for a specific client ID
func getLimiter(clientID string) *rate.Limiter {
	actual, loaded := limiters.Load(clientID)
	if !loaded {
		// Configuration: Refill 1 token per second, max burst capacity of 3 tokens
		limiterConfig := rate.NewLimiter(rate.Every(1*time.Second), 3)
		actual, _ = limiters.LoadOrStore(clientID, &clientLimiter{
			limiter:  limiterConfig,
			lastSeen: time.Now(),
		})
	}

	c := actual.(*clientLimiter)
	
	c.mu.Lock()
	c.lastSeen = time.Now() // Safely update the last-seen timestamp
	c.mu.Unlock()

	return c.limiter
}

// RateLimitMiddleware acts as our primary perimeter checkpoint
func RateLimitMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// In a production ecosystem, resolve this using a User ID extracted from a JWT token.
		// For simplicity, we fallback to using Client IP tracking here.
		clientID := c.ClientIP() 

		limiter := getLimiter(clientID)

		// Check if a token can be consumed immediately (Allow is entirely non-blocking)
		if !limiter.Allow() {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"error":       "Too Many Requests",
				"message":     "You are querying the AI system too frequently. Please slow down.",
				"retry_after": "1s",
			})
			c.Abort() // Halt the pipeline, preventing execution from reaching the core handler
			return
		}

		c.Next()
	}
}

Main Orchestration: main.go

Go

func main() {
	r := gin.Default()

	// Apply the middleware selectively to our AI processing route group
	aiGroup := r.Group("/api/v1/ai")
	aiGroup.Use(RateLimitMiddleware())
	
	aiGroup.POST("/chat", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"status": "success",
			"answer": "This is the generated payload response from the internal AI model...",
		})
	})

	r.Run(":8080")
}

🎯 Daily Mission

Spin up this local Go server, open up a secondary Terminal window, and simulate a rapid burst spam attack using this bash curl loop:

Bash

for i in {1..6}; do curl -X POST http://localhost:8080/api/v1/ai/chat; echo ""; done

Challenge Analysis: Observe the terminal logs carefully. At which exact request iteration does our middleware start blocking traffic and throwing back an HTTP 429 Too Many Requests status code?

To take your architecture design skills a step further: If your corporate policy dictates that "Engineers can query the AI 5 times per second, but the Marketing team is capped at 2 times per second," how would you refactor our getLimiter function to gracefully support dynamic parameter configurations based on user groups? Sketch the architecture logic out in your head!

FAQ: Frequently Asked Questions

Why do we need a sync.RWMutex around the lastSeen field if we are already using a thread-safe sync.Map?

While sync.Map protects map operations (reads/writes to keys) from collapsing across concurrent threads, it does not guard the internal values or fields of structural elements nested inside it. If our background cleanup job reads client.lastSeen to gauge eviction while our middleware is concurrently overwriting c.lastSeen = time.Now(), a Data Race condition occurs. Encapsulating the struct with a dedicated Mutex ensures 100% memory safety.

What is the fundamental difference between Allow() and Wait() in the x/time/rate package?

Allow() is purely non-blocking; it returns a boolean value instantly depending on whether a token is available, making it optimal for API HTTP middlewares. Conversely, Wait() blocks execution (puts the current goroutine to sleep) until a token becomes available, which is ideal for background script orchestrations or workers processing third-party APIs where you want to gracefully throttle yourself to avoid hitting external rate limits.

If we scale horizontally across 10 application servers (Distributed Cluster), will this code still work?

No, this localized approach will lose efficiency. Because token bucket states are maintained completely in-memory, states are locked to independent servers. If a user distributes requests across server 1, 2, and 3, their limits are evaluated completely separately. For a distributed system, you should offload the Rate Limiting state to a centralized system like Redis using Lua scripts for atomic operations.

Conclusion & Next Episode (EP.165)

Implementing a Token Bucket-based Rate Limiter middleware is a simple and highly effective line of defense to keep malicious scripts, runaway loops, and rogue spam operations from compromising your AI processing endpoints.

In the next episode (EP.165): While rate limiters are excellent safeguards against resource starvation, what happens when real, valid business usage across an organization genuinely scales past a single server's computing threshold? Capping users out restricts productivity, and your hardware can only scale vertically so far. Next time, we will scale into cluster environments with Load Balancing AI Servers — Distributing heavy processing workloads across multi-node AI clusters. Stay tuned, Gophers!

Follow Superdev Academy on all platforms: