Building a Simple Load Balancer in Go

Load balancing is an essential component in distributed systems, ensuring that incoming requests are evenly distributed across multiple backend servers. In this post, we’ll build a simple, thread-safe load balancer in Go using a round-robin algorithm. We’ll walk through the code step-by-step, explaining how each part works and how you can adapt it to fit your needs.

You can clone the project from my GitHub repository.

What We’ll Build

Our load balancer will:

• Distribute incoming HTTP requests among multiple backend servers.
• Use a round-robin algorithm to ensure even distribution.
• Be thread-safe, using sync.Mutex for safe access to shared resources.
• Allow dynamic addition and removal of backend servers.

Prerequisites

Make sure you have Go 1.16 or later installed on your machine.

Code Overview

Here’s the complete code for our load balancer:

package main

import (
    "fmt"
    "net/http"
    "net/http/httputil"
    "net/url"
    "sync"
)

// Server interface defines the methods that a backend server should implement.
type Server interface {
    Address() string
    IsAlive() bool
    Serve(rw http.ResponseWriter, r *http.Request)
}

// SimpleServer implements the Server interface and represents a single backend server.
type SimpleServer struct {
    address string
    proxy   *httputil.ReverseProxy
}

// Address returns the address of the server.
func (s *SimpleServer) Address() string {
    return s.address
}

// IsAlive returns the health status of the server. Always returns true in this example.
func (s *SimpleServer) IsAlive() bool {
    return true
}

// Serve forwards the request to the backend server using the reverse proxy.
func (s *SimpleServer) Serve(rw http.ResponseWriter, r *http.Request) {
    s.proxy.ServeHTTP(rw, r)
}

// LoadBalancer manages the distribution of requests to multiple backend servers.
type LoadBalancer struct {
    port            string
    roundRobinCount int
    servers         []Server
    mu              sync.Mutex
}

// NewLoadBalancer creates a new LoadBalancer instance.
func NewLoadBalancer(port string, servers []Server) *LoadBalancer {
    return &LoadBalancer{
        port:            port,
        roundRobinCount: 0,
        servers:         servers,
    }
}

// NewSimpleServer creates a new SimpleServer instance.
func NewSimpleServer(address string) *SimpleServer {
    serverURL, err := url.Parse(address)
    if err != nil {
        panic(fmt.Sprintf("Error parsing server URL %s: %v", address, err))
    }

    return &SimpleServer{
        address: address,
        proxy:   httputil.NewSingleHostReverseProxy(serverURL),
    }
}

// getNextAvailableServer returns the next available server using a round-robin algorithm.
func (lb *LoadBalancer) getNextAvailableServer() Server {
    lb.mu.Lock()
    defer lb.mu.Unlock()

    for i := 0; i < len(lb.servers); i++ {
        server := lb.servers[lb.roundRobinCount%len(lb.servers)]
        lb.roundRobinCount++
        if server.IsAlive() {
            return server
        }
    }

    if len(lb.servers) > 0 {
        return lb.servers[0]
    }

    return nil
}

// serveProxy forwards the request to the next available backend server.
func (lb *LoadBalancer) serveProxy(rw http.ResponseWriter, r *http.Request) {
    targetServer := lb.getNextAvailableServer()
    if targetServer == nil {
        http.Error(rw, "No available servers", http.StatusServiceUnavailable)
        return
    }

    fmt.Printf("Forwarding request to address %s\n", targetServer.Address())
    targetServer.Serve(rw, r)
}

func main() {
    // List of backend servers.
    servers := []Server{
        NewSimpleServer("https://www.facebook.com"),
        NewSimpleServer("http://www.bing.com"),
        NewSimpleServer("https://www.google.com"),
    }
    lb := NewLoadBalancer("8000", servers)
    handleRedirect := func(rw http.ResponseWriter, r *http.Request) {
        lb.serveProxy(rw, r)
    }
    http.HandleFunc("/", handleRedirect)

    fmt.Printf("Serving requests at 'localhost:%v'\n", lb.port)
    http.ListenAndServe(":"+lb.port, nil)
}

Step-by-Step Explanation

1. Server Interface

The Server interface defines the methods that any backend server should implement:

type Server interface {
    Address() string
    IsAlive() bool
    Serve(rw http.ResponseWriter, r *http.Request)
}

2. SimpleServer Struct

The SimpleServer struct implements the Server interface. It holds the server address and a reverse proxy to forward requests:

type SimpleServer struct {
    address string
    proxy   *httputil.ReverseProxy
}

func (s *SimpleServer) Address() string {
    return s.address
}

func (s *SimpleServer) IsAlive() bool {
    return true
}

func (s *SimpleServer) Serve(rw http.ResponseWriter, r *http.Request) {
    s.proxy.ServeHTTP(rw, r)
}

3. LoadBalancer Struct

The LoadBalancer struct manages the distribution of requests to multiple backend servers. It uses a round-robin algorithm to select the next available server:

type LoadBalancer struct {
    port            string
    roundRobinCount int
    servers         []Server
    mu              sync.Mutex
}

func NewLoadBalancer(port string, servers []Server) *LoadBalancer {
    return &LoadBalancer{
        port:            port,
        roundRobinCount: 0,
        servers:         servers,
    }
}

func (lb *LoadBalancer) getNextAvailableServer() Server {
    lb.mu.Lock()
    defer lb.mu.Unlock()

    for i := 0; i < len(lb.servers); i++ {
        server := lb.servers[lb.roundRobinCount%len(lb.servers)]
        lb.roundRobinCount++
        if server.IsAlive() {
            return server
        }
    }

    if len(lb.servers) > 0 {
        return lb.servers[0]
    }

    return nil
}

func (lb *LoadBalancer) serveProxy(rw http.ResponseWriter, r *http.Request) {
    targetServer := lb.getNextAvailableServer()
    if targetServer == nil {
        http.Error(rw, "No available servers", http.StatusServiceUnavailable)
        return
    }

    fmt.Printf("Forwarding request to address %s\n", targetServer.Address())
    targetServer.Serve(rw, r)
}

4. Main Function

The main function initializes the load balancer with a list of backend servers and starts the HTTP server:

func main() {
    servers := []Server{
        NewSimpleServer("https://www.facebook.com"),
        NewSimpleServer("http://www.bing.com"),
        NewSimpleServer("https://www.google.com"),
    }
    lb := NewLoadBalancer("8000", servers)
    handleRedirect := func(rw http.ResponseWriter, r *http.Request) {
        lb.serveProxy(rw, r)
    }
    http.HandleFunc("/", handleRedirect)

    fmt.Printf("Serving requests at 'localhost:%v'\n", lb.port)
    http.ListenAndServe(":"+lb.port, nil)
}

Adapting the Load Balancer

This basic load balancer can be extended in several ways:

• Health Checks: Implement health checks to periodically check the status of backend servers and mark them as unavailable if they fail.
• Dynamic Server Management: Add endpoints to dynamically add or remove backend servers.
• Load Balancing Algorithms: Implement different load balancing algorithms like least connections, IP hash, etc.

Conclusion

We’ve built a simple load balancer in Go that uses a round-robin algorithm to distribute incoming requests. This example serves as a starting point, and you can extend it to meet more complex requirements. Happy coding!