Building a simple load balancer in Go

Load balancers are crucial in modern software development. If you've ever wondered how requests are distributed across multiple servers, or why certain websites feel faster even during heavy traffic, the answer often lies in efficient load balancing.

In this post, we'll build a simple application load balancer using Round Robin algorithm in Go. The aim of this post is to understand how a load balancer works under the hood, step by step.

What is a Load Balancer?

A load balancer is a system that distributes incoming network traffic across multiple servers. It ensures that no single server bears too much load, preventing bottlenecks and improving the overall user experience. Load balancing approach also ensure that if one server fails, then the traffic can be automatically re-routed to another available server, thus reducing the impact of the failure and increasing availability.

Why do we use Load Balancers?

• High availability: By distributing traffic, load balancers ensure that even if one server fails, traffic can be routed to other healthy servers, making the application more resilient.
• Scalability: Load balancers allow you to scale your system horizontally by adding more servers as traffic increases.
• Efficiency: It maximizes resource utilization by ensuring all servers share the workload equally.

Load balancing algorithms

There are different algorithms and strategies to distribute the traffic:

• Round Robin: One of the simplest methods available. It distributes requests sequentially among the available servers. Once it reaches the last server, it starts again from the beginning.
• Weighted Round Robin: Similar to round robin algorithm except each server is assigned some fixed numerical weighting. This given weight is used to determine the server for routing traffic.
• Least Connections: Routes traffic to the server with the least active connections.
• IP Hashing: Select the server based on the client's IP address.

In this post, we'll focus on implementing a Round Robin load balancer.

What is a Round Robin algorithm?

A round robin algorithm sends each incoming request to the next available server in a circular manner. If server A handles the first request, server B will handle the second, and server C will handle the third. Once all servers have received a request, it starts again from server A.

Now, let's jump into the code and build our load balancer!

Step 1: Define the Load Balancer and Server

type LoadBalancer struct {
    Current int
    Mutex   sync.Mutex
}

We'll first define a simple LoadBalancer struct with a Current field to keep track of which server should handle next request. The Mutex ensures that our code is safe to use concurrently.

Each server we load balance is defined by the Server struct:

type Server struct {
    URL       *url.URL
    IsHealthy bool
    Mutex     sync.Mutex
}

Here, each server has a URL and an IsHealthy flag, which indicates whether the server is available to handle requests.

Step 2: Round Robin Algorithm

The heart of our load balancer is the round robin algorithm. Here's how it works:

func (lb *LoadBalancer) getNextServer(servers []*Server) *Server {
    lb.Mutex.Lock()
    defer lb.Mutex.Unlock()

    for i := 0; i < len(servers); i++ {
        idx := lb.Current % len(servers)
        nextServer := servers[idx]
        lb.Current++

        nextServer.Mutex.Lock()
        isHealthy := nextServer.IsHealthy
        nextServer.Mutex.Unlock()

        if isHealthy {
            return nextServer
        }
    }

    return nil
}

• This method loops through the list of servers in a round robin fashion. If the selected server is healthy, it returns that server to handle the incoming request.
• We are using Mutex to ensure that only one goroutine can access and modify the Current field of the load balancer at a time. This ensures that the round robin algorithm operates correctly when multiple requests are being processed concurrently.
• Each Server also has its own Mutex. When we check the IsHealthy field, we lock the server's Mutex to prevent concurrent access from multiple goroutines.
• Without Mutex locking it is possible that another goroutine could be changing the value which could result in reading an incorrect or inconsistent data.
• We unlock the Mutex as soon as we have updated the Current field or read the IsHealthy field value to keep the critical section as small as possible. In this way, we are using Mutex to avoid any race condition.

Step 3: Configuring the Load Balancer

Our configuration is stored in a config.json file, which contains the server URLs and health check intervals (more on it in below section).

type Config struct {
    Port                string   `json:"port"`
    HealthCheckInterval string   `json:"healthCheckInterval"`
    Servers             []string `json:"servers"`
}

The configuration file might look like this:

{
  "port": ":8080",
  "healthCheckInterval": "2s",
  "servers": [
    "http://localhost:5001",
    "http://localhost:5002",
    "http://localhost:5003",
    "http://localhost:5004",
    "http://localhost:5005"
  ]
}

Step 4: Health Checks

We want to make sure that the servers are healthy before routing any incoming traffic to them. This is done by sending periodic health checks to each server:

func healthCheck(s *Server, healthCheckInterval time.Duration) {
    for range time.Tick(healthCheckInterval) {
        res, err := http.Head(s.URL.String())
        s.Mutex.Lock()
        if err != nil || res.StatusCode != http.StatusOK {
            fmt.Printf("%s is down\n", s.URL)
            s.IsHealthy = false
        } else {
            s.IsHealthy = true
        }
        s.Mutex.Unlock()
    }
}

Every few seconds (as specified in the config), the load balancer sends a HEAD request to each server to check if it is healthy. If a server is down, the IsHealthy flag is set to false, preventing future traffic from being routed to it.

Step 5: Reverse Proxy

When the load balancer receives a request, it forwards the request to the next available server using a reverse proxy. In Golang, the httputil package provides a built-in way to handle reverse proxying, and we will use it in our code through the ReverseProxy function:

func (s *Server) ReverseProxy() *httputil.ReverseProxy {
    return httputil.NewSingleHostReverseProxy(s.URL)
}

What is a Reverse Proxy?

A reverse proxy is a server that sits between a client and one or more backend severs. It receives the client's request, forwards it to one of the backend servers, and then returns the server's response to the client. The client interacts with the proxy, unaware of which specific backend server is handling the request.

In our case, the load balancer acts as a reverse proxy, sitting in front of multiple servers and distributing incoming HTTP requests across them.

Step 6: Handling Requests

When a client makes a request to the load balancer, it selects the next available healthy server using the round robin algorithm implementation in getNextServer function and proxies the client request to that server. If no healthy server is available then we send service unavailable error to the client.

http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        server := lb.getNextServer(servers)
        if server == nil {
            http.Error(w, "No healthy server available", http.StatusServiceUnavailable)
            return
        }
        w.Header().Add("X-Forwarded-Server", server.URL.String())
        server.ReverseProxy().ServeHTTP(w, r)
    })

The ReverseProxy method proxies the request to the actual server, and we also add a custom header X-Forwarded-Server for debugging purposes (though in production, we should avoid exposing internal server details like this).

Step 7: Starting the Load Balancer

Finally, we start the load balancer on the specified port:

log.Println("Starting load balancer on port", config.Port)
err = http.ListenAndServe(config.Port, nil)
if err != nil {
        log.Fatalf("Error starting load balancer: %s\n", err.Error())
}

Working Demo

TL;DR

In this post, we built a basic load balancer from scratch in Golang using a round robin algorithm. This is a simple yet effective way to distribute traffic across multiple servers and ensure that your system can handle higher loads efficiently.

There's a lot more to explore, such as adding sophisticated health checks, implementing different load balancing algorithms, or improving fault tolerance. But this basic example can be a solid foundation to build upon.

You can find the source code in this GitHub repo.