Skip to content

Instantly share code, notes, and snippets.

@rednafi
Last active September 26, 2023 17:34
Show Gist options
  • Save rednafi/4f871286f42177f21a74a0ce038ce725 to your computer and use it in GitHub Desktop.
Save rednafi/4f871286f42177f21a74a0ce038ce725 to your computer and use it in GitHub Desktop.
Dummy load balancer in a single Go script. Here's the full explanation: https://rednafi.com/go/dummy_load_balancer
/*
cc Redowan Delowar (rednafi.com)
+----------------------------------------+
| Load Balancer (8080) |
| +----------------------------------+ |
| | | |
| | Request from Client | |
| | | |
| +-----------------|----------------+ |
| | Forward Request |
| | to Backend |
| v |
| +----------------------------------+ |
| | | |
| | Load Balancing | |
| | | |
| | +----------+ +----------+ | |
| | | Backend | | Backend | | |
| | | 8081 | | 8082 | | |
| | +----------+ +----------+ | |
| | | |
| +-----------------|----------------+ |
| | Distribute Load |
| v |
| +----------------------------------+ |
| | | |
| | Backend Servers | |
| | | |
| | +----------+ +----------+ | |
| | | Response | | Response | | |
| | | Body | | Body | | |
| | +----------+ +----------+ | |
| | | |
| +----------------------------------+ |
| | Send Response |
| v |
| +----------------------------------+ |
| | Client receives Response | |
| +----------------------------------+ |
+----------------------------------------+
*/
package main
import (
"fmt"
"io"
"net/http"
"sync"
)
var (
backends = []string{
"http://localhost:8081/b8081",
"http://localhost:8082/b8082",
}
currentBackend int
backendMutex sync.Mutex
)
// Start a backend server on the specified port
func startBackend(port int, wg *sync.WaitGroup) {
// Signals the lb when a backend is done processing a request
defer wg.Done()
http.HandleFunc(fmt.Sprintf("/b%d", port),
func(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "Hello from backend server on :%d\n", port)
})
addr := fmt.Sprintf(":%d", port)
fmt.Printf("Backend is listening on :%d \n", port)
err := http.ListenAndServe(addr, nil)
if err != nil {
fmt.Printf("Error for server on :%d; %s\n", port, err)
}
}
// Get the next backend server to forward the request to
// in a round-robin fashion. This function is thread-safe
func getNextBackend() string {
backendMutex.Lock()
defer backendMutex.Unlock()
backend := backends[currentBackend]
currentBackend = (currentBackend + 1) % len(backends)
return backend
}
// Handle incoming requests and forward them to the backend
func loadBalancerHandler(w http.ResponseWriter, r *http.Request) {
// Pick a backend in round-robin fashion
backend := getNextBackend()
// Relay the client's request to the backend
resp, err := http.Get(backend)
if err != nil {
http.Error(w, "Backend Error", http.StatusInternalServerError)
return
}
defer resp.Body.Close()
// Copy the backend response headers and propagate them to the client
for key, values := range resp.Header {
for _, value := range values {
w.Header().Set(key, value)
}
}
// Copy the backend response body and propagate it to the client
io.Copy(w, resp.Body)
}
func main() {
var wg sync.WaitGroup
ports := []int{8081, 8082}
// Starts the backend servers in the background
for _, port := range ports {
wg.Add(1)
go startBackend(port, &wg)
}
// Starts the load balancer server in the foreground
http.HandleFunc("/", loadBalancerHandler)
fmt.Println("Load balancer is listening on :8080")
err := http.ListenAndServe(":8080", nil)
if err != nil {
fmt.Printf("Error: %s\n", err)
}
}
@rednafi
Copy link
Author

rednafi commented Sep 26, 2023

@JekRock Ah, the blocking call on line 131 is working as a proxy wg.Wait(). Thanks for the detailed explanation and the suggestions.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment