← Reverse Proxy / Load Balancer networkingmedium ready

Problem

Reverse Proxy / Load Balancer (Rust)

Topics: round-robin, health checks, atomics, Service

Problem

Distribute requests across backend servers using round-robin, skipping unhealthy backends. Backends are Services (Request -> Response), so the balancer routes to them in-process — the same hand-written HTTP shape as the other networking challenges, no real sockets.

pub struct Backend { pub url: String, /* health + service */ }
pub struct LoadBalancer { /* … */ }
impl LoadBalancer {
    pub fn new(backends: Vec<Backend>) -> Self;
    pub fn next(&self) -> Result<&Backend, NoHealthyBackends>;
    pub fn set_healthy(&self, url: &str, healthy: bool);
    pub fn health_check(&self, cancel: &Cancel, interval: Duration, path: &str);
}
impl Service for LoadBalancer { /* proxy to next(), or 502 */ }

next returns the next healthy backend in round-robin order; Err(NoHealthyBackends) if none. Try each backend at most once so an all-down pool can't loop forever.
LoadBalancer::call proxies the request to next()'s backend, or responds 502 if none are healthy.
set_healthy flips a backend's health by url.
health_check probes each backend (GET path) immediately and then every interval until cancel fires, marking healthy iff the probe returns 200.

You implement next and the Service::call routing; set_healthy, check_once, and the health_check loop are provided.

Key concepts

Atomics: an AtomicU64 counter for round-robin and an AtomicBool per backend for health — no locks on the hot path.
Round-robin with skip: advance counter % len, but loop at most len times so an all-unhealthy pool returns the error instead of spinning.
Health-check loop: probe on an interval, cancellable via the token — the same shape as a real gateway's background health checker.

Grading

Your solution.rs is compiled together with a trusted tests.rs (which include!s it) using rustc --test. Only solution.rs is yours to edit.

Sign in to submit your solution.

Leaderboard score: lower is better, 1.0 ≈ reference

#	Name	score	ns/op	allocs

Output

No submission yet.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// A cooperative cancellation token (stands in for Go's `context`).
#[derive(Clone, Default)]
pub struct Cancel(Arc<AtomicBool>);

impl Cancel {
    pub fn new() -> Self {
        Cancel(Arc::new(AtomicBool::new(false)))
    }
    pub fn cancel(&self) {
        self.0.store(true, Ordering::SeqCst);
    }
    pub fn is_cancelled(&self) -> bool {
        self.0.load(Ordering::SeqCst)
    }
}

#[derive(Clone, Default)]
pub struct Request {
    pub method: String,
    pub path: String,
}

impl Request {
    pub fn new(method: &str, path: &str) -> Self {
        Request {
            method: method.to_string(),
            path: path.to_string(),
        }
    }
}

pub struct Response {
    pub status: u16,
    pub body: String,
}

impl Response {
    pub fn new(status: u16) -> Self {
        Response {
            status,
            body: String::new(),
        }
    }
    pub fn with_body(mut self, body: &str) -> Self {
        self.body = body.to_string();
        self
    }
}

/// Turns a `Request` into a `Response`. Backends are `Service`s, so the load
/// balancer routes to them in-process (no real sockets).
pub trait Service {
    fn call(&self, req: Request) -> Response;
}

pub struct HandlerFn<F>(pub F);

impl<F: Fn(Request) -> Response> Service for HandlerFn<F> {
    fn call(&self, req: Request) -> Response {
        (self.0)(req)
    }
}

/// Returned by `next` when every backend is unhealthy.
#[derive(Debug, PartialEq, Eq)]
pub struct NoHealthyBackends;

/// One upstream: an identifier, a health flag, and the service that handles its
/// requests (including health probes).
pub struct Backend {
    pub url: String,
    healthy: AtomicBool,
    service: Box<dyn Service + Send + Sync>,
}

impl Backend {
    pub fn new(url: &str, service: impl Service + Send + Sync + 'static) -> Self {
        Backend {
            url: url.to_string(),
            healthy: AtomicBool::new(true),
            service: Box::new(service),
        }
    }
    pub fn healthy(&self) -> bool {
        self.healthy.load(Ordering::SeqCst)
    }
    /// Sends a request to this backend's service.
    pub fn serve(&self, req: Request) -> Response {
        self.service.call(req)
    }
}

/// Round-robin load balancer that skips unhealthy backends.
pub struct LoadBalancer {
    backends: Vec<Backend>,
    counter: AtomicU64,
}

impl LoadBalancer {
    pub fn new(backends: Vec<Backend>) -> Self {
        LoadBalancer {
            backends,
            counter: AtomicU64::new(0),
        }
    }

/// Returns the next healthy backend in round-robin order, trying each at most
    /// once so an all-unhealthy pool can't loop forever. Use `self.counter` and
    /// `Backend::healthy`.
    pub fn next(&self) -> Result<&Backend, NoHealthyBackends> {
        let _ = (&self.backends, &self.counter);
        unimplemented!()
    }

pub fn set_healthy(&self, url: &str, healthy: bool) {
        for b in &self.backends {
            if b.url == url {
                b.healthy.store(healthy, Ordering::SeqCst);
                return;
            }
        }
    }

/// Probes every backend once: GET `path`, mark healthy iff it returns 200.
    pub fn check_once(&self, path: &str) {
        for b in &self.backends {
            let resp = b.serve(Request::new("GET", path));
            b.healthy.store(resp.status == 200, Ordering::SeqCst);
        }
    }

/// Probes immediately, then every `interval` until `cancel` fires.
    pub fn health_check(&self, cancel: &Cancel, interval: Duration, path: &str) {
        self.check_once(path);
        while !cancel.is_cancelled() {
            let deadline = Instant::now() + interval;
            while Instant::now() < deadline {
                if cancel.is_cancelled() {
                    return;
                }
                let remaining = deadline.saturating_duration_since(Instant::now());
                std::thread::sleep(remaining.min(Duration::from_millis(5)));
            }
            if cancel.is_cancelled() {
                return;
            }
            self.check_once(path);
        }
    }
}

impl Service for LoadBalancer {
    /// Proxies the request to the next healthy backend, or responds 502 if none.
    fn call(&self, req: Request) -> Response {
        let _ = req;
        unimplemented!()
    }
}