API Gateway, Service Discovery, and Load Balancing

Split services need three things: a single entry point for outside traffic, a way to find the addresses of each service, and a way to pick one instance per request. This lesson builds a small gateway in Axum with all three.

API gateway: one entry point that routes external requests to internal services and applies shared policy (authentication, limits).
Backend-for-frontend (BFF): a gateway endpoint shaped for one client type. It calls several services and returns one response.
Service discovery: maps a service name to the addresses of its running instances.
Load balancing: picks one instance for each request.
Reverse proxy: forwards a request to another server and returns its response.

Project setup

Cargo.toml

TOML

[package]
name = "storefront-gateway"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = "0.8"
tokio = { version = "1", features = ["full"] }
reqwest = { version = "0.12", features = ["json"] }
serde_json = "1"

The project has two binaries: the gateway (src/main.rs) and a demo upstream service (src/bin/upstream.rs) that runs as several instances.

Example 1: A demo upstream service

The upstream reports its own name, so responses show which instance answered. It stands in for the catalog and inventory services.

src/bin/upstream.rs

Rust

use axum::{
    extract::{Path, State},
    routing::get,
    Json, Router,
};
use serde_json::{json, Value};

async fn product(State(name): State<String>, Path(id): Path<u64>) -> Json<Value> {
    Json(json!({
        "instance": name,
        "id": id,
        "name": format!("Product {id}"),
        "price_cents": 1999
    }))
}

async fn inventory(State(name): State<String>, Path(id): Path<u64>) -> Json<Value> {
    Json(json!({ "instance": name, "id": id, "in_stock": 12 }))
}

#[tokio::main]
async fn main() {
    let mut args = std::env::args().skip(1);
    let name = args.next().expect("usage: upstream <name> <port>");
    let port = args.next().expect("usage: upstream <name> <port>");

    let app = Router::new()
        .route("/api/products/{id}", get(product))
        .route("/api/inventory/{id}", get(inventory))
        .route("/health", get(|| async { "ok" }))
        .with_state(name);

    let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}"))
        .await
        .unwrap();
    axum::serve(listener, app).await.unwrap();
}

Example 2: A reverse proxy with a routing table

A gateway is a routing table (path prefix to upstream) plus shared policy. The proxy function rebuilds the request for the upstream and returns the upstream response.

src/proxy.rs

Rust

use axum::{
    body::{to_bytes, Body},
    extract::Request,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
};
use reqwest::Client;

const HOP_BY_HOP: [&str; 6] = [
    "host",
    "connection",
    "keep-alive",
    "transfer-encoding",
    "upgrade",
    "te",
];

fn strip(headers: &mut HeaderMap) {
    for name in HOP_BY_HOP {
        headers.remove(name);
    }
}

pub async fn forward(client: &Client, base: &str, req: Request) -> Response {
    let (parts, body) = req.into_parts();
    let path = parts
        .uri
        .path_and_query()
        .map(|p| p.as_str())
        .unwrap_or("/");
    let url = format!("{base}{path}");

    let Ok(bytes) = to_bytes(body, 1024 * 1024).await else {
        return StatusCode::PAYLOAD_TOO_LARGE.into_response();
    };

    let mut headers = parts.headers;
    strip(&mut headers);

    let result = client
        .request(parts.method, url)
        .headers(headers)
        .body(bytes)
        .send()
        .await;
    let upstream = match result {
        Ok(response) => response,
        Err(error) if error.is_timeout() => return StatusCode::GATEWAY_TIMEOUT.into_response(),
        Err(_) => return StatusCode::BAD_GATEWAY.into_response(),
    };

    let status = upstream.status();
    let mut response_headers = upstream.headers().clone();
    strip(&mut response_headers);
    let Ok(bytes) = upstream.bytes().await else {
        return StatusCode::BAD_GATEWAY.into_response();
    };

    let mut response = Response::new(Body::from(bytes));
    *response.status_mut() = status;
    *response.headers_mut() = response_headers;
    response
}

Host is removed so the client library sets the upstream host. Hop-by-hop headers describe one connection and are not forwarded.
The body is buffered up to 1 MiB, which protects gateway memory. Streaming avoids buffering for large uploads.
A connection failure returns 502 Bad Gateway. A client timeout returns 504 Gateway Timeout.
A production proxy also sets x-forwarded-for from the connection address.

The first gateway has one upstream per route.

src/main.rs

Rust

mod proxy;

use axum::{
    extract::{Request, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    Router,
};
use reqwest::Client;
use std::time::Duration;

struct Route {
    prefix: &'static str,
    upstream: &'static str,
}

const ROUTES: [Route; 2] = [
    Route { prefix: "/api/products", upstream: "http://127.0.0.1:9001" },
    Route { prefix: "/api/inventory", upstream: "http://127.0.0.1:9003" },
];

async fn gateway(State(client): State<Client>, req: Request) -> Response {
    let path = req.uri().path().to_owned();
    match ROUTES.iter().find(|route| path.starts_with(route.prefix)) {
        Some(route) => proxy::forward(&client, route.upstream, req).await,
        None => StatusCode::NOT_FOUND.into_response(),
    }
}

#[tokio::main]
async fn main() {
    let client = Client::builder()
        .timeout(Duration::from_secs(3))
        .build()
        .unwrap();
    let app = Router::new().fallback(gateway).with_state(client);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

The fallback handler receives every request that no explicit route matches. Its first match wins, so longer prefixes go first.

Example 3: Load balancing and discovery

Server-side balancing puts a proxy in front of the instances (this gateway, or a Kubernetes Service). Client-side balancing lets the caller hold the instance list and choose, which removes one hop and needs a discovery source.

Algorithm	Behavior	Suitable when
Round robin	Instances in turn	Requests cost about the same
Least connections	Fewest active requests	Request cost varies
Two random choices	Pick two at random, use the less loaded	Many callers, no shared state
Consistent hashing	Same key, same instance	Instances hold per-key caches

Discovery mechanisms:

DNS: the name resolves to one or more addresses. Caching delays removal of dead instances.
Registry: instances register and send heartbeats, and callers query the registry. Data is fresh, and the registry is one more component to run.
Kubernetes: a Service gives a stable name and virtual address, and the platform tracks ready Pods. A headless Service returns one DNS record per Pod for client-side balancing.

This example adds a round-robin balancer whose instance list can be replaced, and a background task that refreshes it from DNS.

src/balancer.rs

Rust

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

pub struct Balancer {
    instances: RwLock<Vec<String>>,
    next: AtomicUsize,
}

impl Balancer {
    pub fn new(instances: Vec<String>) -> Self {
        Balancer {
            instances: RwLock::new(instances),
            next: AtomicUsize::new(0),
        }
    }

    pub fn set(&self, instances: Vec<String>) {
        *self.instances.write().unwrap() = instances;
    }

    pub fn pick(&self) -> Option<String> {
        let list = self.instances.read().unwrap();
        if list.is_empty() {
            return None;
        }
        let index = self.next.fetch_add(1, Ordering::Relaxed) % list.len();
        Some(list[index].clone())
    }
}

src/discovery.rs

Rust

use crate::balancer::Balancer;
use std::{sync::Arc, time::Duration};
use tokio::net::lookup_host;

pub fn spawn_dns_refresh(balancer: Arc<Balancer>, host: String, port: u16) {
    tokio::spawn(async move {
        let mut tick = tokio::time::interval(Duration::from_secs(5));
        loop {
            tick.tick().await;
            match lookup_host((host.as_str(), port)).await {
                Ok(addresses) => {
                    let mut list: Vec<String> =
                        addresses.map(|address| format!("http://{address}")).collect();
                    list.sort();
                    list.dedup();
                    if !list.is_empty() {
                        balancer.set(list);
                    }
                }
                Err(error) => eprintln!("dns lookup failed: {error}"),
            }
        }
    });
}

An empty or failed lookup keeps the last known list, so a DNS outage does not empty the balancer. This gateway takes its instance lists from environment variables, and CATALOG_DNS switches the catalog to DNS refresh.

src/main.rs

Rust

mod balancer;
mod discovery;
mod proxy;

use axum::{
    extract::{Request, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    Router,
};
use balancer::Balancer;
use reqwest::Client;
use std::{sync::Arc, time::Duration};

struct Route {
    prefix: &'static str,
    balancer: Arc<Balancer>,
}

#[derive(Clone)]
struct AppState {
    client: Client,
    routes: Arc<Vec<Route>>,
}

fn from_env(var: &str, default: &str) -> Arc<Balancer> {
    let value = std::env::var(var).unwrap_or_else(|_| default.to_owned());
    let instances = value.split(',').map(str::to_owned).collect();
    Arc::new(Balancer::new(instances))
}

async fn gateway(State(state): State<AppState>, req: Request) -> Response {
    let path = req.uri().path().to_owned();
    let Some(route) = state.routes.iter().find(|r| path.starts_with(r.prefix)) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    match route.balancer.pick() {
        Some(base) => proxy::forward(&state.client, &base, req).await,
        None => StatusCode::SERVICE_UNAVAILABLE.into_response(),
    }
}

#[tokio::main]
async fn main() {
    let catalog = from_env("CATALOG_UPSTREAMS", "http://127.0.0.1:9001,http://127.0.0.1:9002");
    let inventory = from_env("INVENTORY_UPSTREAMS", "http://127.0.0.1:9003");

    if let Ok(target) = std::env::var("CATALOG_DNS") {
        let (host, port) = target.rsplit_once(':').expect("CATALOG_DNS is host:port");
        discovery::spawn_dns_refresh(catalog.clone(), host.to_owned(), port.parse().unwrap());
    }

    let state = AppState {
        client: Client::builder().timeout(Duration::from_secs(3)).build().unwrap(),
        routes: Arc::new(vec![
            Route { prefix: "/api/products", balancer: catalog },
            Route { prefix: "/api/inventory", balancer: inventory },
        ]),
    };
    let app = Router::new().fallback(gateway).with_state(state);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

The Kubernetes side of discovery is plain YAML. The first Service load-balances connections behind one virtual address. The second is headless (clusterIP: None), so DNS returns one address per ready Pod.

k8s/catalog-services.yaml

YAML

apiVersion: v1
kind: Service
metadata:
  name: catalog
spec:
  selector:
    app: catalog
  ports:
    - port: 8080
      targetPort: 8080
---
apiVersion: v1
kind: Service
metadata:
  name: catalog-instances
spec:
  clusterIP: None
  selector:
    app: catalog
  ports:
    - port: 8080
      targetPort: 8080

A ClusterIP Service balances connections, so one long-lived HTTP/2 or gRPC connection stays on one Pod. Request-level balancing needs a client-side balancer over the headless Service, or a proxy that understands HTTP/2 (production gRPC).

Example 4: A backend-for-frontend endpoint and gateway policy

A BFF endpoint calls independent services concurrently and merges the results. This example adds src/bff.rs and a gateway API-key check. Token validation in a real gateway follows OAuth, JWT and mTLS. The gateway rejects a bad key before any upstream call, and services still check permissions for each action.

src/bff.rs

Rust

use crate::{balancer::Balancer, AppState};
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    Json,
};
use reqwest::Client;
use serde_json::{json, Value};

async fn fetch(client: &Client, balancer: &Balancer, path: &str) -> Result<Value, ()> {
    let base = balancer.pick().ok_or(())?;
    let response = client
        .get(format!("{base}{path}"))
        .send()
        .await
        .map_err(|_| ())?;
    response
        .error_for_status()
        .map_err(|_| ())?
        .json::<Value>()
        .await
        .map_err(|_| ())
}

pub async fn product_page(State(state): State<AppState>, Path(id): Path<u64>) -> Response {
    let (product, stock) = tokio::join!(
        fetch(&state.client, &state.catalog, &format!("/api/products/{id}")),
        fetch(&state.client, &state.inventory, &format!("/api/inventory/{id}")),
    );
    match (product, stock) {
        (Ok(product), Ok(stock)) => Json(json!({ "product": product, "stock": stock })).into_response(),
        _ => StatusCode::BAD_GATEWAY.into_response(),
    }
}

tokio::join! runs both calls at once, so the endpoint takes as long as the slower call. The final main.rs in the next section wires this in. Deadlines, bounded retries (idempotent calls only), and circuit breakers at the gateway follow resilience patterns.

Complete example

Text

storefront-gateway/
  Cargo.toml
  k8s/catalog-services.yaml
  src/
    main.rs
    proxy.rs
    balancer.rs
    discovery.rs
    bff.rs
    bin/upstream.rs

Cargo.toml, src/proxy.rs, src/balancer.rs, src/discovery.rs, src/bin/upstream.rs, and k8s/catalog-services.yaml are the files from Examples 1 to 3, unchanged. They are repeated here so this section runs on its own.

Cargo.toml

TOML

[package]
name = "storefront-gateway"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = "0.8"
tokio = { version = "1", features = ["full"] }
reqwest = { version = "0.12", features = ["json"] }
serde_json = "1"

src/bin/upstream.rs

Rust

use axum::{
    extract::{Path, State},
    routing::get,
    Json, Router,
};
use serde_json::{json, Value};

async fn product(State(name): State<String>, Path(id): Path<u64>) -> Json<Value> {
    Json(json!({
        "instance": name,
        "id": id,
        "name": format!("Product {id}"),
        "price_cents": 1999
    }))
}

async fn inventory(State(name): State<String>, Path(id): Path<u64>) -> Json<Value> {
    Json(json!({ "instance": name, "id": id, "in_stock": 12 }))
}

#[tokio::main]
async fn main() {
    let mut args = std::env::args().skip(1);
    let name = args.next().expect("usage: upstream <name> <port>");
    let port = args.next().expect("usage: upstream <name> <port>");

    let app = Router::new()
        .route("/api/products/{id}", get(product))
        .route("/api/inventory/{id}", get(inventory))
        .route("/health", get(|| async { "ok" }))
        .with_state(name);

    let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}"))
        .await
        .unwrap();
    axum::serve(listener, app).await.unwrap();
}

src/proxy.rs

Rust

use axum::{
    body::{to_bytes, Body},
    extract::Request,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
};
use reqwest::Client;

const HOP_BY_HOP: [&str; 6] = [
    "host",
    "connection",
    "keep-alive",
    "transfer-encoding",
    "upgrade",
    "te",
];

fn strip(headers: &mut HeaderMap) {
    for name in HOP_BY_HOP {
        headers.remove(name);
    }
}

pub async fn forward(client: &Client, base: &str, req: Request) -> Response {
    let (parts, body) = req.into_parts();
    let path = parts
        .uri
        .path_and_query()
        .map(|p| p.as_str())
        .unwrap_or("/");
    let url = format!("{base}{path}");

    let Ok(bytes) = to_bytes(body, 1024 * 1024).await else {
        return StatusCode::PAYLOAD_TOO_LARGE.into_response();
    };

    let mut headers = parts.headers;
    strip(&mut headers);

    let result = client
        .request(parts.method, url)
        .headers(headers)
        .body(bytes)
        .send()
        .await;
    let upstream = match result {
        Ok(response) => response,
        Err(error) if error.is_timeout() => return StatusCode::GATEWAY_TIMEOUT.into_response(),
        Err(_) => return StatusCode::BAD_GATEWAY.into_response(),
    };

    let status = upstream.status();
    let mut response_headers = upstream.headers().clone();
    strip(&mut response_headers);
    let Ok(bytes) = upstream.bytes().await else {
        return StatusCode::BAD_GATEWAY.into_response();
    };

    let mut response = Response::new(Body::from(bytes));
    *response.status_mut() = status;
    *response.headers_mut() = response_headers;
    response
}

src/balancer.rs

Rust

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

pub struct Balancer {
    instances: RwLock<Vec<String>>,
    next: AtomicUsize,
}

impl Balancer {
    pub fn new(instances: Vec<String>) -> Self {
        Balancer {
            instances: RwLock::new(instances),
            next: AtomicUsize::new(0),
        }
    }

    pub fn set(&self, instances: Vec<String>) {
        *self.instances.write().unwrap() = instances;
    }

    pub fn pick(&self) -> Option<String> {
        let list = self.instances.read().unwrap();
        if list.is_empty() {
            return None;
        }
        let index = self.next.fetch_add(1, Ordering::Relaxed) % list.len();
        Some(list[index].clone())
    }
}

src/discovery.rs

Rust

use crate::balancer::Balancer;
use std::{sync::Arc, time::Duration};
use tokio::net::lookup_host;

pub fn spawn_dns_refresh(balancer: Arc<Balancer>, host: String, port: u16) {
    tokio::spawn(async move {
        let mut tick = tokio::time::interval(Duration::from_secs(5));
        loop {
            tick.tick().await;
            match lookup_host((host.as_str(), port)).await {
                Ok(addresses) => {
                    let mut list: Vec<String> =
                        addresses.map(|address| format!("http://{address}")).collect();
                    list.sort();
                    list.dedup();
                    if !list.is_empty() {
                        balancer.set(list);
                    }
                }
                Err(error) => eprintln!("dns lookup failed: {error}"),
            }
        }
    });
}

src/bff.rs

Rust

use crate::{balancer::Balancer, AppState};
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    Json,
};
use reqwest::Client;
use serde_json::{json, Value};

async fn fetch(client: &Client, balancer: &Balancer, path: &str) -> Result<Value, ()> {
    let base = balancer.pick().ok_or(())?;
    let response = client
        .get(format!("{base}{path}"))
        .send()
        .await
        .map_err(|_| ())?;
    response
        .error_for_status()
        .map_err(|_| ())?
        .json::<Value>()
        .await
        .map_err(|_| ())
}

pub async fn product_page(State(state): State<AppState>, Path(id): Path<u64>) -> Response {
    let (product, stock) = tokio::join!(
        fetch(&state.client, &state.catalog, &format!("/api/products/{id}")),
        fetch(&state.client, &state.inventory, &format!("/api/inventory/{id}")),
    );
    match (product, stock) {
        (Ok(product), Ok(stock)) => Json(json!({ "product": product, "stock": stock })).into_response(),
        _ => StatusCode::BAD_GATEWAY.into_response(),
    }
}

src/main.rs

Rust

mod balancer;
mod bff;
mod discovery;
mod proxy;

use axum::{
    extract::{Request, State},
    http::StatusCode,
    middleware::{from_fn, Next},
    response::{IntoResponse, Response},
    routing::get,
    Router,
};
use balancer::Balancer;
use reqwest::Client;
use std::{sync::Arc, time::Duration};

#[derive(Clone)]
struct AppState {
    client: Client,
    catalog: Arc<Balancer>,
    inventory: Arc<Balancer>,
}

fn from_env(var: &str, default: &str) -> Arc<Balancer> {
    let value = std::env::var(var).unwrap_or_else(|_| default.to_owned());
    let instances = value.split(',').map(str::to_owned).collect();
    Arc::new(Balancer::new(instances))
}

async fn require_key(req: Request, next: Next) -> Response {
    let valid = req
        .headers()
        .get("x-api-key")
        .is_some_and(|value| value == "storefront-key");
    if valid {
        next.run(req).await
    } else {
        StatusCode::UNAUTHORIZED.into_response()
    }
}

async fn gateway(State(state): State<AppState>, req: Request) -> Response {
    let path = req.uri().path().to_owned();
    let balancer = if path.starts_with("/api/products") {
        &state.catalog
    } else if path.starts_with("/api/inventory") {
        &state.inventory
    } else {
        return StatusCode::NOT_FOUND.into_response();
    };
    match balancer.pick() {
        Some(base) => proxy::forward(&state.client, &base, req).await,
        None => StatusCode::SERVICE_UNAVAILABLE.into_response(),
    }
}

#[tokio::main]
async fn main() {
    let catalog = from_env("CATALOG_UPSTREAMS", "http://127.0.0.1:9001,http://127.0.0.1:9002");
    let inventory = from_env("INVENTORY_UPSTREAMS", "http://127.0.0.1:9003");

    if let Ok(target) = std::env::var("CATALOG_DNS") {
        let (host, port) = target.rsplit_once(':').expect("CATALOG_DNS is host:port");
        discovery::spawn_dns_refresh(catalog.clone(), host.to_owned(), port.parse().unwrap());
    }

    let state = AppState {
        client: Client::builder().timeout(Duration::from_secs(3)).build().unwrap(),
        catalog,
        inventory,
    };

    let protected = Router::new()
        .route("/bff/product-page/{id}", get(bff::product_page))
        .fallback(gateway)
        .layer(from_fn(require_key));

    let app = Router::new()
        .route("/health", get(|| async { "ok" }))
        .merge(protected)
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

require_key guards the BFF route and the proxy fallback. /health stays open for probes.
The single /api/ routing table became two balancers on AppState so the BFF and the proxy share them.

k8s/catalog-services.yaml is the YAML shown in Example 3. It is not used by the local run.

Run (four terminals):

Terminal

cargo run --bin upstream -- catalog-a 9001
cargo run --bin upstream -- catalog-b 9002
cargo run --bin upstream -- inventory-a 9003
cargo run --bin storefront-gateway

Requests, in this order (the balancer counter starts at zero):

Terminal

curl -s -H 'x-api-key: storefront-key' http://127.0.0.1:8080/api/products/7
curl -s -H 'x-api-key: storefront-key' http://127.0.0.1:8080/api/products/7
curl -s -H 'x-api-key: storefront-key' http://127.0.0.1:8080/bff/product-page/7
curl -s -o /dev/null -w '%{http_code}\n' http://127.0.0.1:8080/api/products/7
curl -s -o /dev/null -w '%{http_code}\n' -H 'x-api-key: storefront-key' http://127.0.0.1:8080/nothing
curl -s http://127.0.0.1:8080/health

Output

{"id":7,"instance":"catalog-a","name":"Product 7","price_cents":1999}
{"id":7,"instance":"catalog-b","name":"Product 7","price_cents":1999}
{"product":{"id":7,"instance":"catalog-a","name":"Product 7","price_cents":1999},"stock":{"id":7,"in_stock":12,"instance":"inventory-a"}}
401
404
ok

The first two responses alternate between instances. The third request is the third catalog pick, so catalog-a answers again. Stopping catalog-b makes about half of the catalog calls return 502 until discovery removes it, which is why health-based removal and retries are part of production discovery.

Summary

A gateway is a routing table plus shared edge policy. Business rules stay in services.
A reverse proxy rebuilds the request, strips hop-by-hop headers, limits body size, and maps failures to 502 and 504.
A BFF endpoint calls services concurrently with tokio::join! and returns one shaped response.
Discovery maps a service name to live instances through DNS, a registry with heartbeats, or the Kubernetes platform.
A Kubernetes ClusterIP Service balances connections. HTTP/2 and gRPC traffic need request-level balancing.
Client-side balancing removes a hop and needs an instance list. Server-side balancing needs no client support.
Balancing algorithms include round robin, least connections, two random choices, and consistent hashing.
Retries are safe only for idempotent operations and need a cap.