Building High-Throughput Microservices with Rust and Axum: Complete Production Guide

.
├── Cargo.toml
├── src
│   ├── main.rs
│   ├── config.rs
│   ├── handlers.rs
│   ├── models.rs
│   ├── middleware
│   │   ├── mod.rs
│   │   ├── rate_limit.rs
│   │   └── tracing.rs
│   └── db.rs
└── Dockerfile

# Cargo.toml
[package]
name = "high-throughput-service"
version = "0.1.0"
edition = "2021"

[dependencies]
# Web framework
axum = { version = "0.7", features = ["macros"] }
tokio = { version = "1", features = ["full"] }
tower = { version = "0.4", features = ["full"] }
tower-http = { version = "0.5", features = ["full"] }

# Database
sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "postgres", "uuid", "chrono"] }
deadpool-redis = { version = "0.13", features = ["rt-tokio-1"] } # For rate limiting

# Serialization/Deserialization
serde = { version = "1", features = ["derive"] }
serde_json = "1"

# Configuration
config = { version = "0.13", features = ["toml", "yaml"] } # Or `envy` for env vars

# Logging and Tracing
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
opentelemetry = { version = "0.21", features = ["rt-tokio"] }
opentelemetry-sdk = { version = "0.21", features = ["rt-tokio"] }
opentelemetry-stdout = { version = "0.16" } # For local dev
opentelemetry-otlp = { version = "0.14", features = ["grpc-tonic", "reqwest-client", "tokio"] }
tracing-opentelemetry = "0.22"

# Utilities
uuid = { version = "1", features = ["v4", "serde"] }
chrono = { version = "0.4", features = ["serde"] }
anyhow = "1"

// src/config.rs
use serde::Deserialize;
use std::net::SocketAddr;

#[derive(Debug, Deserialize, Clone)]
pub struct AppConfig {
    pub server_address: SocketAddr,
    pub database_url: String,
    pub redis_url: String,
    pub service_name: String,
    pub otel_exporter_otlp_endpoint: Option<String>,
}

impl AppConfig {
    pub fn load() -> Result<Self, config::ConfigError> {
        let s = config::Config::builder()
            .add_source(config::File::with_name("config.toml").required(false))
            .add_source(config::Environment::with_prefix("APP"))
            .build()?;
        s.try_deserialize()
    }
}

// src/db.rs
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use deadpool_redis::{Pool as RedisPool, Config as RedisConfig, Runtime};
use std::time::Duration;
use anyhow::Result;

pub async fn setup_pg_pool(database_url: &str) -> Result<PgPool> {
    let pool = PgPoolOptions::new()
        .max_connections(50)
        .min_connections(10)
        .acquire_timeout(Duration::from_secs(5))
        .connect_timeout(Duration::from_secs(5))
        .idle_timeout(Duration::from_secs(30 * 60)) // 30 minutes
        .test_before_acquire(true)
        .connect(database_url)
        .await?;
    Ok(pool)
}

pub fn setup_redis_pool(redis_url: &str) -> Result<RedisPool> {
    let cfg = RedisConfig::from_url(redis_url);
    let pool = cfg.create_pool(Some(Runtime::Tokio1))?;
    Ok(pool)
}

// src/middleware/tracing.rs
use axum::{
    extract::Request,
    middleware::Next,
    response::Response,
};
use opentelemetry::{
    global,
    trace::{SpanKind, Tracer},
    Context, KeyValue,
};
use opentelemetry_sdk::{
    trace::{Config, TracerProvider},
    Resource,
};
use opentelemetry_otlp::WithExportConfig;
use tracing::{info_span, Instrument};
use tracing_opentelemetry::OpenTelemetrySpanExt;
use anyhow::Result;

pub fn init_tracer(service_name: &str, otlp_endpoint: Option<&str>) -> Result<()> {
    let resource = Resource::new(vec![
        KeyValue::new("service.name", service_name.to_string()),
        KeyValue::new("service.version", env!("CARGO_PKG_VERSION").to_string()),
    ]);

    let tracer_provider = if let Some(endpoint) = otlp_endpoint {
        opentelemetry_otlp::new_exporter()
            .tonic()
            .with_endpoint(endpoint)
            .build_tracer_provider_with_config(
                Config::default().with_resource(resource)
            )
    } else {
        // Fallback to stdout for local development if no OTLP endpoint is configured
        opentelemetry_stdout::new_pipeline()
            .with_trace_config(Config::default().with_resource(resource))
            .install_simple()
            .expect("Failed to install stdout tracer provider")
    };

    global::set_tracer_provider(tracer_provider);
    Ok(())
}

pub async fn trace_layer(req: Request, next: Next) -> Response {
    let path = req.uri().path().to_string();
    let method = req.method().to_string();

    let tracer = global::tracer("axum-server");
    let parent_cx = global::get_text_map_propagator(|propagator| {
        propagator.extract(&opentelemetry_http::HeaderExtractor(req.headers()))
    });

    let span = tracer
        .span_builder(format!("HTTP {} {}", method, path))
        .with_kind(SpanKind::Server)
        .with_parent_context(parent_cx)
        .start(&tracer);

    let cx = Context::current_with_span(span);
    let _guard = cx.attach();

    let response = next.run(req).instrument(info_span!("request_processing")).await;

    let span = cx.span();
    span.set_attribute(KeyValue::new("http.method", method));
    span.set_attribute(KeyValue::new("http.target", path));
    span.set_attribute(KeyValue::new("http.status_code", response.status().as_u16() as i64));
    span.end();

    response
}

// src/middleware/rate_limit.rs
use axum::{
    extract::{Request, State},
    http::StatusCode,
    middleware::Next,
    response::Response,
};
use deadpool_redis::redis::AsyncCommands;
use deadpool_redis::Pool as RedisPool;
use std::time::Duration;
use tracing::{error, info};

#[derive(Clone)]
pub struct RateLimiter {
    pub redis_pool: RedisPool,
    pub limit: u64,
    pub window_seconds: u64,
}

impl RateLimiter {
    pub fn new(redis_pool: RedisPool, limit: u64, window_seconds: u64) -> Self {
        Self {
            redis_pool,
            limit,
            window_seconds,
        }
    }

    pub async fn layer(State(rate_limiter): State<RateLimiter>, request: Request, next: Next) -> Result<Response, StatusCode> {
        let ip_address = request
            .headers()
            .get("X-Forwarded-For")
            .and_then(|h| h.to_str().ok())
            .and_then(|s| s.split(',').next()) // Take the first IP in case of multiple proxies
            .unwrap_or("unknown")
            .to_string();

        let key = format!("rate_limit:{}", ip_address);

        let mut conn = rate_limiter.redis_pool.get().await.map_err(|e| {
            error!("Failed to get Redis connection: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

        let (count, _): (u64, u64) = deadpool_redis::redis::pipe()
            .atomic()
            .incr(&key, 1)
            .expire(&key, rate_limiter.window_seconds as usize)
            .query_async(&mut *conn)
            .await
            .map_err(|e| {
                error!("Failed to execute Redis rate limit command: {:?}", e);
                StatusCode::INTERNAL_SERVER_ERROR
            })?;

        if count > rate_limiter.limit {
            info!("Rate limit exceeded for IP: {}", ip_address);
            return Err(StatusCode::TOO_MANY_REQUESTS);
        }

        Ok(next.run(request).await)
    }
}

// src/handlers.rs
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::{IntoResponse, Json},
};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use uuid::Uuid;
use chrono::{DateTime, Utc};
use tracing::{info, error};

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Item {
    pub id: Uuid,
    pub name: String,
    pub description: Option<String>,
    pub created_at: DateTime<Utc>,
}

#[derive(Debug, Deserialize)]
pub struct CreateItem {
    pub name: String,
    pub description: Option<String>,
}

#[derive(Clone)]
pub struct AppState {
    pub pg_pool: PgPool,
    // Other shared state like RedisPool, etc.
}

pub async fn health_check() -> impl IntoResponse {
    (StatusCode::OK, "OK")
}

pub async fn create_item(
    State(state): State<AppState>,
    Json(payload): Json<CreateItem>,
) -> Result<Json<Item>, StatusCode> {
    info!("Attempting to create item: {}", payload.name);
    let new_item = sqlx::query_as!(
        Item,
        r#"
        INSERT INTO items (id, name, description, created_at)
        VALUES ($1, $2, $3, $4)
        RETURNING id, name, description, created_at
        "#,
        Uuid::new_v4(),
        payload.name,
        payload.description,
        Utc::now()
    )
    .fetch_one(&state.pg_pool)
    .await
    .map_err(|e| {
        error!("Failed to insert item: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    info!("Successfully created item with ID: {}", new_item.id);
    Ok(Json(new_item))
}

pub async fn get_item(
    State(state): State<AppState>,
    Path(item_id): Path<Uuid>,
) -> Result<Json<Item>, StatusCode> {
    info!("Attempting to retrieve item with ID: {}", item_id);
    let item = sqlx::query_as!(
        Item,
        r#"
        SELECT id, name, description, created_at
        FROM items
        WHERE id = $1
        "#,
        item_id
    )
    .fetch_optional(&state.pg_pool)
    .await
    .map_err(|e| {
        error!("Failed to query item: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?
    .ok_or(StatusCode::NOT_FOUND)?;

    info!("Successfully retrieved item with ID: {}", item_id);
    Ok(Json(item))
}

// src/main.rs (excerpt)
// ... imports ...
use tokio::signal;
use tracing::info;

async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    info!("Shutdown signal received, initiating graceful shutdown...");
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // ... config loading, tracing setup, db setup ...

    let app = Router::new()
        .route("/health", get(health_check))
        .route("/items", post(create_item))
        .route("/items/:id", get(get_item))
        .with_state(app_state.clone())
        .layer(middleware::from_fn_with_state(rate_limiter.clone(), RateLimiter::layer))
        .layer(middleware::from_fn(trace_layer))
        .layer(TraceLayer::new_for_http()) // Axum's built-in tracing for request/response logging
        .layer(SetRequestIdLayer::new(
            Make<Uuid>::new(),
            PropagateRequestIdLayer::new(HeaderName::from_static("x-request-id")),
        ))
        .layer(TimeoutLayer::new(Duration::from_secs(30))); // Request timeout

    let listener = tokio::net::TcpListener::bind(&config.server_address).await?;
    info!("Listening on {}", config.server_address);

    axum::serve(listener, app.into_make_service())
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    info!("Server gracefully shut down.");
    Ok(())
}

# Dockerfile

# Stage 1: Builder
FROM rust:1.78-slim-bookworm AS builder

# Install build dependencies for SQLx and OpenSSL
RUN apt-get update && apt-get install -y \
    pkg-config \
    libssl-dev \
    postgresql-client \
    musl-tools \
    && rm -rf /var/lib/apt/lists/*

# Set the target for musl (static linking)
RUN rustup target add x86_64-unknown-linux-musl

WORKDIR /app

# Copy Cargo.toml and Cargo.lock first to leverage Docker cache
COPY Cargo.toml Cargo.lock ./

# Create a dummy src directory and main.rs to cache dependencies
RUN mkdir src && echo "fn main() {}" > src/main.rs
# Build dependencies only
RUN cargo build --release --target x86_64-unknown-linux-musl
# Remove dummy files
RUN rm -rf src target/x86_64-unknown-linux-musl/release/high-throughput-service

# Copy the actual source code
COPY src ./src

# Build the application
RUN cargo build --release --target x86_64-unknown-linux-musl

# Stage 2: Runtime
FROM scratch

# Set timezone data (important for many applications)
# FROM debian:bookworm-slim # Alternative if scratch is too restrictive for your needs
# RUN apt-get update && apt-get install -y tzdata && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Copy the compiled binary from the builder stage
COPY --from=builder /app/target/x86_64-unknown-linux-musl/release/high-throughput-service ./high-throughput-service

# Copy configuration file if present
COPY config.toml ./config.toml

# Expose the port the application listens on
EXPOSE 8080

# Set environment variables for configuration
ENV APP_SERVER_ADDRESS="0.0.0.0:8080"
ENV APP_DATABASE_URL="postgres://user:password@host:port/database"
ENV APP_REDIS_URL="redis://:password@host:port/"
ENV APP_SERVICE_NAME="high-throughput-service"
# ENV APP_OTEL_EXPORTER_OTLP_ENDPOINT="http://otel-collector:4317"

# Run the application
CMD ["./high-throughput-service"]

// Example: Custom validator extractor
use axum::{
    async_trait,
    extract::{FromRequest, Request},
    http::StatusCode,
    response::{IntoResponse, Response},
    Json,
};
use serde::de::DeserializeOwned;
use validator::Validate;

pub struct ValidatedJson<T>(pub T);

#[async_trait]
impl<T, S> FromRequest<S> for ValidatedJson<T>
where
    T: DeserializeOwned + Validate,
    S: Send + Sync,
{
    type Rejection = Response;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let Json(value) = Json::<T>::from_request(req, state)
            .await
            .map_err(|err| err.into_response())?;

        value.validate().map_err(|err| {
            (StatusCode::BAD_REQUEST, Json(err)).into_response()
        })?;

        Ok(ValidatedJson(value))
    }
}

// Usage in handler:
// pub async fn create_item(ValidatedJson(payload): ValidatedJson<CreateItem>) -> ...

// src/main.rs (excerpt for migrations)
// ...
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // ...
    let pg_pool = db::setup_pg_pool(&config.database_url).await?;

    // Apply database migrations
    info!("Running database migrations...");
    sqlx::migrate!("./migrations") // Path to your migrations directory
        .run(&pg_pool)
        .await?;
    info!("Database migrations applied successfully.");
    // ...
}