Aagaman Luitel

Writing a load balancer in golang

Writing a load balancer in golang

2024-10-22 | 14mins

Table Of Contents

Reverse proxy

From cloudfare: A reverse proxy is a server that sits in front of web servers and forwards client (e.g. web browser) requests to those web servers.
So our objective is to create a server that feeds in the request, send the request to one of our server and send back the request. And to make it more interesting let us make our reverse proxy decide the efficient server. Our reverse proxy will sit on Layer 4 of OSI model (Transport layer) and choose server it wants. This is essentially a L4 load balancer.

Existing load balancer

I tried https://github.com/yyyar/gobetween which is a L4 load balancer with a configuration system. All you have to do is define your servers in gobetween.toml file and run the load balancer. And your request to loadbalancer will be forwarded into servers. Now that we know what we want lets simulate our load balancer

  • Loadbalancer takes an request

  • Server1 is located on http://localhost:4444

  • Server2 is located on http://localhost:4445

  • Our load balancer (reverse proxy) is located on http://localhost:3000

    • It has 2 servers configured at localhost:{4444,4445}

  • request come to http://localhost:3000

  • forwards the request to Server1 or Server2

  • get the result

Writing our own

We will be using golang since it just works and thats what I'm most comfortable when it comes to developing servers.

Server endpoints

Lets define our server, instead of using configuration files and parsing them, I will be hardcoding them into our source.

const (
	SERVER1 = "localhost:4444"
	SERVER2 = "localhost:4445"
	SERVER3 = "localhost:4446"
)

Creating dummy servers

Make a separate folder and write the dummy servers. We will not be writing any super complicated servers but they should work as well. I have tested this load balancer in some of my local http projects and it just works but for now let us be simple. This server just prints which server it is, how many times it has been called and the headers it is receiving.

package main

import (
	"fmt"
	"net/http"
)

var v int = 0

func HomePage(w http.ResponseWriter, r *http.Request) {
	v++
	fmt.Printf("Server1 called %d times\n", v)
	fmt.Println(w.Header(), r.Header)
	w.Write([]byte("hello world from server 1"))
}

func main() {
	http.HandleFunc("/", HomePage)
	err := http.ListenAndServe("localhost:4444", nil)
	if err != nil {
		panic(err)
	}
}

Create 3-4 of these and put the address in SERVERS in our loadbalancer.

The load balancer

Golang has a powerful standard library which lets us use net, http, templates, crypto and many more without external library. For now we will only be using net since it provides TCP/UDP network methods which we need.

Our loadbalancer is still a server even if it makes request to another server.

Listen announces on the local network address. The network must be "tcp", "tcp4", "tcp6", "unix" or "unixpacket". Sounds like what we want. A server with tcp capabilities.

func main(){
	listener, err := net.Listen("tcp", "localhost:3000")
	if err != nil {
		panic(err)
	}

Note that localhost does not makes our server accessible from other machines, we need to use 0.0.0.0. This applies everywhere. Use "localhost:3000" if you are communicating in your machine only or "0.0.0.0:3000" if communicating through IPs. I will be using "localhost:3000" since I am not making request through another device/ip.
And now we want to actually listen using the listener

...
		panic(err)
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}

Accept waits for and returns the next connection to the listener. It returns a new connection from the incoming request. Lets do a simple ping test

...
		panic(err)
	}

	for {
		_, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		fmt.Println("conn")
	}

Terminal (main)

 owl@archlinux~ $ go run main.go

Terminal (curl)

 owl@archlinux~ $ curl localhost:3000

Terminal (main)

 owl@archlinux~ $ go run main.go
 conn
 conn

Seems to be working.
net also provides methods to connect to other servers with net.Dial

	for {
		_, err := listener.Accept()
		if err != nil {
			panic(err)
		}

		_, err = net.Dial("tcp", SERVER1)
		if err != nil {
			panic(err)
		}
	}

Dial one of our server as soon as we get a request. This program panics if SERVER1 is not active. Definitely a good thing. Now we read the incoming request and send to one of our servers.

	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		// read from our l4 request

		connDial, err := net.Dial("tcp", SERVER1)
		if err != nil {
			panic(err)
		}
		// send the response to the server
		// read the response from the server
		// send the response to l4

		fmt.Printf("sent\n%s\nto server: %s\n", recvBuf, SERVER1)
	}

Now replacing comments with code we get a reverse proxy in our hands. This works becase http is essentially a TCP connection with http headers. When receiving http headers on our TCP server we will just forward that http headers to the underlying servers. This is not the same for https servers where it has a TLS/SSL layer on top which makes the header encrypted.


	sendBuf := make([]byte, 1024)
	recvBuf := make([]byte, 1024)
	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		conn.Read(sendBuf)

		connDial, err := net.Dial("tcp", SERVER1)
		if err != nil {
			panic(err)
		}
		connDial.Write(sendBuf)

		connDial.Read(recvBuf)

		conn.Write(recvBuf)

		fmt.Printf("sent\n%s\nto server: %s\n\nReceived %s", sendBuf, recvBuf, SERVER1)
	}

The actual load balancer

Now that we know the basics of our reverse proxy, lets add a load balancer algorithm and decide the server on runtime. This can be implemented however we like. I'll be focusing on 2 algorithms. One is silly. Before lets make above code readable and reusable.


func main() {
	listener, err := net.Listen("tcp", "localhost:3000")
	if err != nil {
		panic(err)
	}

	// buffer to store the fields we get from the incoming request
	recvBuf := make([]byte, 1024)
	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		fmt.Printf("accepting a request...\n\n")
		conn.Read(recvBuf)

		b := dial(recvBuf, SERVER1)
		conn.Write(b)
	}
}

// we dial one of our server with http fields stored in fields buffer
func dial(fields []byte, address string) []byte {
	fmt.Printf("dialing %s\n", address)
	conn, err := net.Dial("tcp", address)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	recvBuf := make([]byte, 1024)

	conn.Write(fields)
	_, err = conn.Read(recvBuf)
	if err != nil {
		panic(err)
	}
	return recvBuf
}

Round Robin Algorithm

Round robin algorithm distributes work evenly among all available resources. In our case we will send request to SERVER1, SERVER2,... SERVER Nth, SERVER1, SERVER2.... and so on where nth is the last server. This is not a full proof algorithm since in real world there are better algorithms suited where the load balancer has the workload stats of the server and it chooses the most optimal one based on resource availability. We use round robin because its one of the algorithm used to create load balancer, and its simple.

Fortunately go makes this super easy. Lets create a struct called stateful servers since every server will have a state where the state is number of time it is called

var StatefulServers = map[string]int{
	SERVER1: 0,
	SERVER2: 0,
	SERVER3: 0,
}

Now for round robin we want our dial function to call any SERVER. Any SERVER that gets returned, its state is increased by one. For this we will find out the SERVER with least state. And return that server.

func RoundRobin(fields []byte) []byte {
	// we need to get lowest value from our stateful servers
	prev := StatefulServers[SERVER1]

  // find least state
	for _, v := range StatefulServers {
		if v < prev {
			prev = v
		}
	}

	var server string
	for k, v := range StatefulServers {
		if v == prev {
			server = k
		}
	}

	StatefulServers[server]++

	return dial(fields, server)
}

And thats it! Our load balancer is now complete.


func main(){
...
		fmt.Printf("accepting a request...\n\n")
		conn.Read(recvBuf)

		b := RoundRobin(recvBuf)
		conn.Write(b)
		}
}

Random Robin

This is just choosing random servers but since I used round robin earlier I will call this random robin. Kind of silly but this is also the fun in programming.

var servers = []string{
	SERVER1,
	SERVER2,
	SERVER3,
}

func RandomRobin(fields []byte) []byte {
	i := rand.Int() % len(servers)
	return dial(fields, servers[i])
}

Use RandomRobin or RoundRobin the fact that our reverse proxy chooses a server on its own makes it a load balancer that just works.

Accepting HTTPS connections

According to cloudfare: HTTPS is just the HTTP protocol but with data encryption using SSL/TLS. So adding a tls config to our loadbalancer should make our http only to https supported. Here is a great gist on how to create certificate and use tls in golang. And I followed the same guide and will be referencing this as the guide.

  • https://gist.github.com/denji/12b3a568f092ab951456

Modifying our original server to https

After creating certificates

...
func main() {
	http.HandleFunc("/", HomePage)
  // err := http.ListenAndServe("localhost:4444", nil)
	err := http.ListenAndServeTLS(":4446", "server.crt", "server.key", nil)
	if err != nil {
		panic(err)
	}
...

Upgrade loadbalancer to https (optional)

A load balancer can stay without tls and can make request to https connection. Its not a requirement for communication but we will do here since why not. Again, make a new certificate or use your old ones we used in our original server. For tls config we should replace net with tls "crypto/package". LoadX509KeyPair parses a public and private key from a PEM encoded files (our openssl certificate).

func main() {
	cer, err := tls.LoadX509KeyPair("server.crt", "server.key")
	if err != nil {
		panic(err)
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cer},
	}
	listener, err := tls.Listen("tcp", "localhost:3000", tlsConfig)
	if err != nil {
		panic(err)
	}
	...

This is server part, now our loadbalancer can be called through https://localhost:3000 but it still cannot communicate to https servers.

Loadbalancer to support https request/response

For our load balancer to make a tls request, the connection should include the tls certificate of the server. In our case we are assuming every connection is https and have same tls config of server.crt. This cal also be seen in config file of gobetween. https://github.com/yyyar/gobetween/blob/6e185295c8476810c64a27f5da28edd778e23423/config/gobetween.toml#L153

...
func dial(fields []byte, address string) []byte {
	cert, err := os.ReadFile("server.crt")
	if err != nil {
		panic(err)
	}
	certPool := x509.NewCertPool()
	certPool.AppendCertsFromPEM(cert)

	conf := &tls.Config{
		RootCAs: certPool,
	}

	fmt.Printf("dialing %s\n", address)
	conn, err := tls.Dial("tcp", address, conf)
	if err != nil {
		panic(err)
}
...

If we don't want to verify the tls of servers we can just ignore the verification process, this way we don't have to remember tls certificate of every server but we are essentially ignoring the tls step. It should only be used for testing.

...
func dial(fields []byte, address string) []byte {
    conf := &tls.Config{
      InsecureSkipVerify: true,
    }

    conn, err := tls.Dial("tcp", address, conf)
    if err != nil {
...

Finalizing

Now all that is left is to create a proper configuration system. I will be coding these in the code itself, creating a struct for server like this

type Server struct {
	address string
	tlspath string
	https   bool
	state   int
}

The other code is pretty much same with servers defined in main and changing data types of our Round Robin algorithm. Here is the final code for our load balancer that support http as well as https

package main

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"os"
)

type Server struct {
	address string
	tlspath string
	https   bool
	state   int
}

const (
	SERVER1 = "localhost:4444"
	SERVER2 = "localhost:4445"
	SERVER3 = "localhost:4446"
)

func main() {
	servers := []Server{
		{
			address: SERVER1,
			https:   true,
			state:   1,
		},
		{
			address: SERVER2,
			https:   false,
			state:   1,
		},
		{
			address: SERVER3,
			https:   true,
			state:   1,
		},
	}

	cer, err := tls.LoadX509KeyPair("server.crt", "server.key")
	if err != nil {
		panic(err)
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cer},
	}
	listener, err := tls.Listen("tcp", "localhost:3000", tlsConfig)
	if err != nil {
		panic(err)
	}

	recvBuf := make([]byte, 1024)
	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		fmt.Printf("accepting a request...\n\n")
		conn.Read(recvBuf)

		b := RoundRobin(recvBuf, servers)
		conn.Write(b)
	}
}

func dial(fields []byte, address string, istls bool) []byte {
	fmt.Printf("dialing %s\n", address)
	recvBuf := make([]byte, 1024)

	if istls {
		cert, err := os.ReadFile("server.crt")
		if err != nil {
			panic(err)
		}
		certPool := x509.NewCertPool()
		certPool.AppendCertsFromPEM(cert)

		conf := &tls.Config{
			RootCAs: certPool,
		}

		conn, err := tls.Dial("tcp", address, conf)
		if err != nil {
			panic(err)
		}
		defer conn.Close()

		conn.Write(fields)
		_, err = conn.Read(recvBuf)
		if err != nil {
			panic(err)
		}
	} else {
		conn, err := net.Dial("tcp", address)
		if err != nil {
			panic(err)
		}
		defer conn.Close()

		conn.Write(fields)
		_, err = conn.Read(recvBuf)
		if err != nil {
			panic(err)
		}
	}

	return recvBuf
}

func (s *Server) IncrementState() {
	s.state += 1
}

func RoundRobin(fields []byte, servers []Server) []byte {
	lowest := servers[0].state

	for _, v := range servers {
		if v.state < lowest {
			lowest = v.state
		}
	}

	var server string
	istls := false
	for i, v := range servers {
		if v.state == lowest {
			server = v.address
			servers[i].IncrementState()
			istls = servers[i].https
			break
		}
	}

	return dial(fields, server, istls)
}