Skip to Content
Course content

227: Building a Simple HTTP Load Balancer in Rust

Click on the "Edit" button in the top corner of the screen to edit your slide content.

I've always found that the best way to understand how a load balancer actually works is to stop reading diagrams and just try to move a request from point A to point B. We aren't building Nginx today, but we are going to build a basic proxy that distributes traffic across a few backend servers. I've got two simple HTTP servers running locally on ports 8081 and 8082—just basic "I am server A" and "I am server B" responses.

Just getting the traffic to move

My first instinct is usually to keep it dead simple. I want to listen on port 8080 and, whenever a request comes in, just forward it to one of my backends. I'll use tokio for the runtime and reqwest as the client to handle the forwarding. Here is my first attempt at the core logic:

async fn handle_request(req: Request) -> Response {
    let backend_url = "http://127.0.0.1:8081";
    let client = reqwest::Client::new();
    
    let res = client.get(backend_url).send().await.unwrap();
    // ... convert reqwest response to our server response ...
}

I ran this, and it worked. Great. But then I noticed something obvious: it's not a load balancer; it's just a proxy. Every single request is hitting server 8081. I'm not actually balancing anything. If 8081 gets slammed, 8082 is just sitting there idling, which defeats the whole purpose of having a cluster.

Wait, it's only hitting one server

Okay, so I need a list of servers and a way to cycle through them. The most straightforward approach is Round Robin. I'll put my server URLs in a vector and keep track of an index. Every time a request hits, I increment the index and move to the next server.

I tried adding a simple usize counter outside the handler, but as soon as I tried to increment it inside the async function, the compiler started yelling at me. Rust doesn't let me mutate a variable from multiple asynchronous tasks without some synchronization. I can't just have a mut index: usize floating around because tokio spawns these handlers across multiple threads.

I could use a Mutex, but for a simple counter, a Mutex feels like overkill. It's a bit heavy for just incrementing a number. Instead, I'll use an AtomicUsize. It's thread-safe, lock-free, and exactly what this scenario calls for.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

struct LoadBalancer {
    backends: Vec<String>,
    current: Arc<AtomicUsize>,
}

impl LoadBalancer {
    fn next_backend(&self) -> &String {
        let index = self.current.fetch_add(1, Ordering::SeqCst);
        &self.backends[index % self.backends.len()]
    }
}

Now, when I hit the endpoint, I see the responses alternating: "Server A", "Server B", "Server A", "Server B". It's working. But there's a catch. I'm creating a new reqwest::Client inside every single request handler. I quickly realized that this is a performance nightmare. Each client creation allocates a connection pool; if I'm doing this 1,000 times a second, I'm going to exhaust my available sockets and crash the OS.

Moving the client to shared state

I need to treat the reqwest::Client as a long-lived object. Since reqwest::Client internally uses an Arc, I can just wrap my entire LoadBalancer state in an Arc and pass it into my route handlers. This way, every single request uses the same connection pool, and the AtomicUsize keeps the balancing logic consistent across threads.

Let's look at how this looks when we put it all together. I'm using axum here because it plays nicely with tokio and makes the routing logic clean:

use axum::{extract::State, routing::get, Router};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

struct AppState {
    client: reqwest::Client,
    backends: Vec<String>,
    counter: AtomicUsize,
}

async fn proxy_handler(State(state): State<Arc<AppState>>) -> String {
    // Round robin selection
    let idx = state.counter.fetch_add(1, Ordering::SeqCst);
    let target = &state.backends[idx % state.backends.len()];
    
    println!("Routing request to: {}", target);
    
    let res = state.client.get(target).send().await;
    
    match res {
        Ok(resp) => resp.text().await.unwrap_or_else(|_| "Error reading body".to_string()),
        Err(_) => "Backend server unavailable".to_string(),
    }
}

#[tokio::main]
async fn main() {
    let state = Arc::new(AppState {
        client: reqwest::Client::new(),
        backends: vec![
            "http://127.0.0.1:8081".to_string(),
            "http://127.0.0.1:8082".to_string(),
        ],
        counter: AtomicUsize::new(0),
    });

    let app = Router::new().route("/", get(proxy_handler)).with_state(state);
    
    let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

This is significantly more robust. I've avoided the "too many open files" error by sharing the client, and the AtomicUsize ensures that the load is distributed evenly without the overhead of a heavy lock. It's simple, it's fast, and it demonstrates the core loop of a load balancer: Intercept $\rightarrow$ Select $\rightarrow$ Forward $\rightarrow$ Return.




📋 Practical Task

Implementing a Passive Health Check

Right now, the load balancer is "blind." If one of your backend servers crashes, the load balancer will still try to send 50% of the traffic to it, resulting in a "Backend server unavailable" message for half of your users. This is unacceptable in a production environment.

Your Task: Modify the proxy_handler to implement a basic "fail-over" mechanism. If the selected backend returns an error (meaning the server is down), the load balancer should immediately try the next available backend in the list before giving up and returning an error to the user.

  • Update the proxy_handler logic to loop through the backends.
  • Ensure that you don't enter an infinite loop if all backends are down.
  • Keep using the AtomicUsize to ensure the starting point for each request is still balanced.
Rating
0 0

There are no comments for now.

to be the first to leave a comment.