Building a Load Balancer in Go: A Practical Guide

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.