Building a Load Balancer in Go: A Practical Guide
Introduction
Load balancing is a critical component in distributed systems, distributing incoming traffic across multiple backend servers to ensure high availability and reliability. While there are many off-the-shelf solutions like NGINX or HAProxy, building a custom load balancer in Go gives you full control over routing logic, health checks, and performance tuning. In this article, we'll implement a simple but functional load balancer using Go's concurrency primitives and standard library.
Core Design
Our load balancer will support:
- Round-robin request distribution
- Passive health checks (marking servers as down after failures)
- Graceful connection handling
- Thread-safe backend pool updates
We'll use Go's net/http/httputil package for reverse proxying and sync for safe concurrent access.
Backend Pool
First, we define a Backend struct to hold server details and a Pool to manage them:
package main
import (
"net/url"
"sync"
"sync/atomic"
)
type Backend struct {
URL *url.URL
Alive bool
mux sync.RWMutex
ReverseProxy *httputil.ReverseProxy
}
func (b *Backend) SetAlive(alive bool) {
b.mux.Lock()
b.Alive = alive
b.mux.Unlock()
}
func (b *Backend) IsAlive() bool {
b.mux.RLock()
alive := b.Alive
b.mux.RUnlock()
return alive
}
type ServerPool struct {
backends []*Backend
current uint64
mux sync.RWMutex
}
func (s *ServerPool) NextIndex() int {
return int(atomic.AddUint64(&s.current, 1) % uint64(len(s.backends)))
}
func (s *ServerPool) GetNextPeer() *Backend {
next := s.NextIndex()
l := len(s.backends) + next
for i := next; i < l; i++ {
idx := i % len(s.backends)
if s.backends[idx].IsAlive() {
if i != next {
atomic.StoreUint64(&s.current, uint64(idx))
}
return s.backends[idx]
}
}
return nil
}
func (s *ServerPool) AddBackend(backend *Backend) {
s.mux.Lock()
s.backends = append(s.backends, backend)
s.mux.Unlock()
}
func (s *ServerPool) HealthCheck() {
for _, b := range s.backends {
alive := isBackendAlive(b.URL.String())
b.SetAlive(alive)
}
}
Health Check
A simple TCP dial health check:
func isBackendAlive(u string) bool {
parsedURL, err := url.Parse(u)
if err != nil {
return false
}
conn, err := net.DialTimeout("tcp", parsedURL.Host, 2*time.Second)
if err != nil {
return false
}
defer conn.Close()
return true
}
Load Balancer Handler
We create an HTTP handler that proxies requests to the next alive backend:
func lbHandler(w http.ResponseWriter, r *http.Request) {
peer := serverPool.GetNextPeer()
if peer != nil {
peer.ReverseProxy.ServeHTTP(w, r)
return
}
http.Error(w, "Service not available", http.StatusServiceUnavailable)
}
Graceful Shutdown
To handle graceful shutdown, we listen for OS signals and drain connections:
func main() {
// ... setup backends and server pool ...
srv := &http.Server{
Addr: ":8080",
Handler: http.HandlerFunc(lbHandler),
}
go func() {
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
log.Fatalf("listen: %s\n", err)
}
}()
quit := make(chan os.Signal, 1)
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
<-quit
log.Println("Shutting down server...")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := srv.Shutdown(ctx); err != nil {
log.Fatalf("Server forced to shutdown: %v", err)
}
log.Println("Server exiting")
}
Putting It All Together
Here's the complete main function with backend initialization:
var serverPool ServerPool
func main() {
backends := []string{
"http://localhost:8081",
"http://localhost:8082",
"http://localhost:8083",
}
for _, backendURL := range backends {
u, _ := url.Parse(backendURL)
proxy := httputil.NewSingleHostReverseProxy(u)
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
log.Printf("[%s] %s\n", u.Host, err.Error())
retry(w, r)
}
serverPool.AddBackend(&Backend{
URL: u,
Alive: true,
ReverseProxy: proxy,
})
}
// Start health check loop
go func() {
t := time.NewTicker(10 * time.Second)
for range t.C {
serverPool.HealthCheck()
}
}()
// ... server setup and graceful shutdown as above ...
}
func retry(w http.ResponseWriter, r *http.Request) {
peer := serverPool.GetNextPeer()
if peer != nil {
peer.ReverseProxy.ServeHTTP(w, r)
} else {
http.Error(w, "Service not available", http.StatusServiceUnavailable)
}
}
Testing
Start three simple HTTP servers (e.g., using Python) and run the load balancer. Requests to localhost:8080 will be distributed evenly across the backends. If one backend goes down, health check will mark it as dead and requests will be routed to the remaining ones.
Conclusion
Building a load balancer in Go is straightforward thanks to its powerful standard library and concurrency model. This implementation provides a solid foundation that can be extended with weighted routing, sticky sessions, or more sophisticated health checks. The full code is available on GitHub.