This 150-Line Go Script Is Actually a Full-On Load Balancer

If you're running services that should be able to handle a high amount of traffic, you can load balance a lot of that traffic between your backend servers. There are lots of production-grade load balancer on the market (NGINX, HAProxy, etc.) but knowing how they work behind the scene is good knowledge.

A Simple HTTP Load Balancer in Go Using Standard Library, In this implementation, we'll use a round-robin algorithm to evenly distribute incoming requests among a collection of backend servers.

The Basic Structure

First, we need to define our core data structures. Our load balancer will track multiple backend servers and their health:

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"sync"
	"sync/atomic"
	"time"
)

// Backend represents a backend server
type Backend struct {
	URL          *url.URL
	Alive        bool
	mux          sync.RWMutex
	ReverseProxy *httputil.ReverseProxy
}

// SetAlive updates the alive status of backend
func (b *Backend) SetAlive(alive bool) {
	b.mux.Lock()
	b.Alive = alive
	b.mux.Unlock()
}

// IsAlive returns true when backend is alive
func (b *Backend) IsAlive() (alive bool) {
	b.mux.RLock()
	alive = b.Alive
	b.mux.RUnlock()
	return
}

// LoadBalancer represents a load balancer
type LoadBalancer struct {
	backends []*Backend
	current  uint64
}

// NextBackend returns the next available backend to handle the request
func (lb *LoadBalancer) NextBackend() *Backend {
	// Simple round-robin
	next := atomic.AddUint64(&lb.current, uint64(1)) % uint64(len(lb.backends))
	
	// Find the next available backend
	for i := 0; i < len(lb.backends); i++ {
		idx := (int(next) + i) % len(lb.backends)
		if lb.backends[idx].IsAlive() {
			return lb.backends[idx]
		}
	}
	return nil
}

Now a few key points here to keep in mind:

1. The Backend struct represents a single backend server with its URL and alive status.
2. We're using a mutex to safely update and check the alive status of each backend in a concurrent environment.
3. The LoadBalancer keeps track of a list of backends and maintains a counter for the round-robin algorithm.
4. We use atomic operations to safely increment the counter in a concurrent environment.
5. The NextBackend method implements the round-robin algorithm, skipping unhealthy backends.

Health Checking

Detecting unavailability of the backend is one of the critical functions of any load balancer. Lets implement a very simple health checking mechanism:

// isBackendAlive checks whether a backend is alive by establishing a TCP connection
func isBackendAlive(u *url.URL) bool {
	timeout := 2 * time.Second
	conn, err := net.DialTimeout("tcp", u.Host, timeout)
	if err != nil {
		log.Printf("Site unreachable: %s", err)
		return false
	}
	defer conn.Close()
	return true
}

// HealthCheck pings the backends and updates their status
func (lb *LoadBalancer) HealthCheck() {
	for _, b := range lb.backends {
		status := isBackendAlive(b.URL)
		b.SetAlive(status)
		if status {
			log.Printf("Backend %s is alive", b.URL)
		} else {
			log.Printf("Backend %s is dead", b.URL)
		}
	}
}

// HealthCheckPeriodically runs a routine health check every interval
func (lb *LoadBalancer) HealthCheckPeriodically(interval time.Duration) {
	t := time.NewTicker(interval)
	for {
		select {
		case <-t.C:
			lb.HealthCheck()
		}
	}
}

This health checker is straightforward:

1. We try to initiate a TCP connection with the backend.
2. If this connection succeeds, then the backend is alive; otherwise it is dead.
3. The check is run at an interval specified in the HealthCheckPeriodically function.

In a production environment, you probably would like a more advanced health check that actually makes an HTTP request to some specified endpoint, but this gets us started.

The HTTP Handler

Let’s go ahead and implement the HTTP handler that will receive request and route them to our backends:

// ServeHTTP implements the http.Handler interface for the LoadBalancer
func (lb *LoadBalancer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	backend := lb.NextBackend()
	if backend == nil {
		http.Error(w, "Service Unavailable", http.StatusServiceUnavailable)
		return
	}
	
	// Forward the request to the backend
	backend.ReverseProxy.ServeHTTP(w, r)
}

And this is where the magic happens:

1. Our round-robin algorithm gives us the next available backend.
2. In the case no backends are available, we return a 503 Service Unavailable error.
3. Otherwise we proxy the request to the selected backend via Go’s internal reverse proxy.

Notice how the net/http/httputil package provides the ReverseProxy type, which handles all the complexities of proxying HTTP requests for us.

Putting It All Together

Finally, let's implement the main function to configure and start our load balancer:

func main() {
	// Parse command line flags
	port := flag.Int("port", 8080, "Port to serve on")
	flag.Parse()

	// Configure backends
	serverList := []string{
		"http://localhost:8081",
		"http://localhost:8082",
		"http://localhost:8083",
	}

	// Create load balancer
	lb := LoadBalancer{}

	// Initialize backends
	for _, serverURL := range serverList {
		url, err := url.Parse(serverURL)
		if err != nil {
			log.Fatal(err)
		}

		proxy := httputil.NewSingleHostReverseProxy(url)
		proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
			log.Printf("Error: %v", err)
			http.Error(w, "Service Unavailable", http.StatusServiceUnavailable)
		}

		lb.backends = append(lb.backends, &Backend{
			URL:          url,
			Alive:        true,
			ReverseProxy: proxy,
		})
		log.Printf("Configured backend: %s", url)
	}

	// Initial health check
	lb.HealthCheck()

	// Start periodic health check
	go lb.HealthCheckPeriodically(time.Minute)

	// Start server
	server := http.Server{
		Addr:    fmt.Sprintf(":%d", *port),
		Handler: &lb,
	}

	log.Printf("Load Balancer started at :%d\n", *port)
	if err := server.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}

Now, here’s what’s going on in the main function:

1. We configure the port based on the command line flags we parse.
2. We set up a list of backend servers.
3. For each backend server, we:
   • Parse the URL
   • Create a reverse proxy for that backend
   • Add error handling for when a backend fails
   • Add the backend to our load balancer
4. We perform an initial health check.
5. We start a goroutine for periodic health checks.
6. Finally, we start the HTTP server with our load balancer as the handler.

Testing the Load Balancer

So again to test our load balancer, we need some backend servers. A naive implementation of a backend server might look like this:

// Save this to backend.go
package main

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

func main() {
	port := flag.Int("port", 8081, "Port to serve on")
	flag.Parse()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		hostname, _ := os.Hostname()
		fmt.Fprintf(w, "Backend server on port %d, host: %s, Request path: %s\n", *port, hostname, r.URL.Path)
	})

	log.Printf("Backend started at :%d\n", *port)
	if err := http.ListenAndServe(fmt.Sprintf(":%d", *port), nil); err != nil {
		log.Fatal(err)
	}
}

Run multiple instances of this backend on different ports:

go build -o backend backend.go
./backend -port 8081 &
./backend -port 8082 &
./backend -port 8083 &

Then build and run the load balancer:

go build -o load-balancer main.go
./load-balancer

Now, make some requests to test it:

curl http://localhost:8080/test

Make multiple requests and you'll see them being distributed across your backend servers in a round-robin fashion.

Potential Improvements

This is a minimal implementation, but there are many things you can do to improve it:

1. Different balancing algorithms: Implement weighted round-robin, least connections, or IP hash-based selection.

2. Better health checking: Make full HTTP requests to a health endpoint instead of just checking TCP connectivity.

3. Metrics collection: Track request counts, response times, and error rates for each backend.

4. Sticky sessions: Ensure requests from the same client always go to the same backend.

5. TLS support: Add HTTPS for secure connections.

6. Dynamic configuration: Allow updating the backend list without restarting.

   
   

We have constructed a simple but effective HTTP load balancer, using nothing but Go’s standard library. This example illustrates important concepts in network programming and the power of Go's built-in networking features.

While this isn't production-ready, it provides a solid foundation for understanding how load balancers work. The complete solution is around 150 lines of code - a testament to the expressiveness of Go and the strength of its standard library.

For production use, you’d want to build out more features and add robust error handling, but the core concepts are all the same. An understanding of these basic principles will prepare you to better configure and debug any load balancer you’ll use down the road.

You can find source code here https://github.com/rezmoss/simple-load-balancer