Building a Rate-Limiter in Go

mkdir rate-limit-proxy
cd rate-limit-proxy
go mod init github.com/yourusername/rate-limit-proxy
touch cmd/rate-limit-proxy/main.go

go get golang.org/x/time/rate

package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// ClientLimiter tracks rate limits for each client
type ClientLimiter struct {
	limiter  *rate.Limiter
	lastSeen time.Time
}

// Global map to store limiters for each client IP
var (
	clients    = make(map[string]*ClientLimiter)
	mu         sync.Mutex // protects clients map
	cleanupInt = 5 * time.Minute
)

func main() {
	// Configuration from environment variables
	targetURL := getEnv("TARGET_URL", "http://localhost:8080")
	proxyPort := getEnv("PROXY_PORT", "8000")
	rpsStr := getEnv("RPS", "10")
	burstStr := getEnv("BURST", "20")

	// Parse numeric configs
	rps, err := strconv.ParseFloat(rpsStr, 64)
	if err != nil {
		log.Fatalf("Invalid RPS value: %v", err)
	}

	burst, err := strconv.Atoi(burstStr)
	if err != nil {
		log.Fatalf("Invalid BURST value: %v", err)
	}

	// Setup the reverse proxy
	url, err := url.Parse(targetURL)
	if err != nil {
		log.Fatalf("Invalid target URL: %v", err)
	}

	proxy := httputil.NewSingleHostReverseProxy(url)

	// Start cleanup routine for inactive clients
	go cleanupOldClients()

	// Set up the HTTP server with our custom handler
	http.HandleFunc("/", createRateLimitHandler(proxy, rps, burst))

	fmt.Printf("Starting rate-limiting proxy server on :%s\n", proxyPort)
	fmt.Printf("Forwarding to: %s\n", targetURL)
	fmt.Printf("Rate limit: %.1f requests/second with burst of %d\n", rps, burst)

	if err := http.ListenAndServe(":"+proxyPort, nil); err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}

// Helper function to get env variable with fallback
func getEnv(key, fallback string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return fallback
}

// Create our rate-limiting handler function
func createRateLimitHandler(proxy *httputil.ReverseProxy, rps float64, burst int) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Get client IP
		clientIP := r.RemoteAddr

		// Get or create rate limiter for this client
		limiter := getClientLimiter(clientIP, rps, burst)

		// Check if request is allowed
		if limiter.limiter.Allow() {
			// Forward the request to the target service
			proxy.ServeHTTP(w, r)
		} else {
			http.Error(w, "429 Too Many Requests", http.StatusTooManyRequests)
			log.Printf("Rate limit exceeded for client %s", clientIP)
		}
	}
}

// Get a rate limiter for a client
func getClientLimiter(clientIP string, rps float64, burst int) *ClientLimiter {
	mu.Lock()
	defer mu.Unlock()

	limiter, exists := clients[clientIP]
	if !exists {
		// Create a new rate limiter for this client
		limiter = &ClientLimiter{
			limiter:  rate.NewLimiter(rate.Limit(rps), burst),
			lastSeen: time.Now(),
		}
		clients[clientIP] = limiter
		return limiter
	}

	// Update the last seen time
	limiter.lastSeen = time.Now()
	return limiter
}

// Cleanup routine to remove inactive clients
func cleanupOldClients() {
	for {
		time.Sleep(cleanupInt)

		mu.Lock()
		now := time.Now()
		for ip, client := range clients {
			if now.Sub(client.lastSeen) > 1*time.Hour {
				delete(clients, ip)
				log.Printf("Cleaned up limiter for inactive client: %s", ip)
			}
		}
		mu.Unlock()
	}
}

// Get client's real IP, respecting X-Forwarded-For header
func getRealIP(r *http.Request) string {
	// Check for X-Forwarded-For header
	forwarded := r.Header.Get("X-Forwarded-For")
	if forwarded != "" {
		// X-Forwarded-For can contain multiple IPs; use the first one
		ips := strings.Split(forwarded, ",")
		return strings.TrimSpace(ips[0])
	}

	// Fall back to RemoteAddr if X-Forwarded-For is not present
	ip := r.RemoteAddr

	// Strip port number if present
	host, _, err := net.SplitHostPort(ip)
	if err == nil {
		// Successfully split host and port, use just the host
		ip = host
	}

	return ip
}

// Create our rate-limiting handler function
func createRateLimitHandler(proxy *httputil.ReverseProxy, rps float64, burst int) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// Get client IP
		clientIP := getRealIP(r)

		// Rest of the function remains the same
		// ...
	}
}

package main

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

func main() {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "Hello from the test service! Path: %s\n", r.URL.Path)
	})
	log.Println("Starting test service on :8080")
	http.ListenAndServe(":8080", nil)
}

# Start a simple test service
go run cmd/echo-server/main.go

export TARGET_URL=http://localhost:8080
export PROXY_PORT=8000
export RPS=2
export BURST=3
go run cmd/rate-limit-proxy/main.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

func main() {
	// Number of concurrent requests to send
	concurrency := 10
	totalRequests := 20

	fmt.Printf("Sending %d requests with %d concurrent workers...\n", totalRequests, concurrency)

	// Use a WaitGroup to wait for all goroutines to finish
	var wg sync.WaitGroup
	wg.Add(concurrency)

	// Channel to distribute work
	jobs := make(chan int, totalRequests)

	// Fill the jobs channel
	for i := 1; i <= totalRequests; i++ {
		jobs <- i
	}
	close(jobs)

	// Start time for benchmarking
	startTime := time.Now()

	// Launch workers
	for i := 0; i < concurrency; i++ {
		go func() {
			defer wg.Done()

			for jobID := range jobs {
				// Make the request
				resp, err := http.Get("http://localhost:8000/")

				if err != nil {
					fmt.Printf("Request %d: ERROR - %v\n", jobID, err)
					continue
				}

				// Read and discard body to free connection
				body, _ := io.ReadAll(resp.Body)
				resp.Body.Close()

				status := "SUCCESS"
				if resp.StatusCode != 200 {
					status = "RATE LIMITED"
				}

				// Print result with timestamp
				elapsed := time.Since(startTime).Seconds()
				fmt.Printf("[%.3fs] Request %d: %s (Status: %d)\n",
					elapsed, jobID, status, resp.StatusCode)

				if resp.StatusCode == 200 {
					// Show a bit of the body for successful requests
					bodyPreview := string(body)
					if len(bodyPreview) > 30 {
						bodyPreview = bodyPreview[:30] + "..."
					}
					fmt.Printf("  Response: %s\n", bodyPreview)
				}
			}
		}()
	}

	// Wait for all workers to finish
	wg.Wait()

	// Show summary
	elapsed := time.Since(startTime)
	fmt.Printf("\nCompleted in %.2f seconds\n", elapsed.Seconds())
	fmt.Printf("Average RPS attempted: %.1f\n", float64(totalRequests)/elapsed.Seconds())
}

export BURST=10
go run cmd/rate-limit-proxy/main.go