Memory Leak in Axum with Dynamodb

A memory leak in an Axum service that uses the AWS SDK for Rust to communicate with DynamoDB typically arises from resource retention across requests. In a long-running Axum application, each incoming HTTP request that performs DynamoDB operations can allocate buffers, SDK clients, or deserialized response structures. If these allocations are held unintentionally—such as by storing request-scientific data in static or long-lived structures, or by failing to drop large query results—the process heap grows over time. This pattern is observable when paginated scans or queries return large items and the application keeps references or caches them beyond the request lifetime.

The Axum runtime does not inherently cause leaks, but its use of async handlers and tower layers can obscure ownership. When a DynamoDB client is cloned per handler rather than shared safely, or when futures capture large payloads by mistake, the Rust runtime may retain memory even after the response is sent. This is especially relevant when integrating middleware or logging that inspects responses containing DynamoDB attribute values, as copies of the payload may linger in buffers or tokio task caches. The leak manifests as gradual RSS growth, increased GC pressure (if using a runtime with GC), and eventual throttling or restarts.

middleBrick scans can surface this risk indirectly by detecting insecure patterns such as missing pagination limits, unencrypted data exposure, or unsafe consumption behaviors that often coexist with inefficient memory handling. For example, an unbounded scan on a DynamoDB table without pagination tokens can return many pages, and if the Axum handler accumulates items into a growing vector, the service becomes susceptible to denial-of-service via resource exhaustion. Proper instrumentation and observability are required to correlate these findings with runtime behavior.

To mitigate memory leaks when using DynamoDB in Axum, structure your handlers to avoid retaining references across requests and to bound data consumption. Use a shared, thread-safe DynamoDB client created once at startup and passed into handlers via Axum extractors. Ensure that pagination is explicitly controlled and that response bodies are streamed or dropped promptly.

Below is a concrete, working Axum example that initializes the AWS SDK client safely and performs a paginated query with bounded result collection:

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_dynamodb::{Client, IntoRequest};
use aws_sdk_dynamodb::model::{AttributeValue, Select};
use std::sync::Arc;
use axum::{routing::get, Router};
use tower_http::trace::TraceLayer;

// Shared client wrapper to avoid cloning the heavy SDK client per request
struct AppState {
    client: Arc,
    table_name: String,
}

async fn query_items(state: axum::extract::State<Arc<AppState>>) -> Result<String, (axum::http::StatusCode, String)> {
    let state = state.0;
    let mut last_evaluated_key: Option<std::collections::HashMap<String, AttributeValue>> = None;
    let mut accumulated_items: Vec<AttributeValue> = Vec::new();

    // Limit page size to prevent unbounded memory growth
    const PAGE_SIZE: i32 = 25;

    loop {
        let mut request = state.client.scan()
            .table_name(&state.table_name)
            .select(Select::AllAttributes)
            .limit(PAGE_SIZE);

        if let Some(ref key) = last_evaluated_key {
            request = request.exclusive_start_key(key.clone());
        }

        let output = request.send().await.map_err(|e| {
            (axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("DynamoDB error: {}", e))
        })?;

        if let Some(items) = output.items() {
            // Important: do not accumulate unbounded items; process in-stream
            for item in items {
                accumulated_items.push(item.clone());
                // Early bound: stop if we reached a safe threshold
                if accumulated_items.len() >= 1000 {
                    return Err((axum::http::StatusCode::REQUEST_ENTITY_TOO Large, "Result set too large".to_string()));
                }
            }
        }

        last_evaluated_key = output.last_evaluated_key().cloned();
        if last_evaluated_key.is_none() {
            break;
        }
    }

    // Here you would serialize accumulated_items safely and return or stream them
    Ok(format!("Retrieved {} items", accumulated_items.len()))
}

#[tokio::main]
async fn main() {
    let region_provider = RegionProviderChain::default_provider().or_else("us-east-1");
    let config = aws_config::from_env().region(region_provider).load().await;
    let client = Arc::new(Client::new(&config));
    let state = Arc::new(AppState {
        client,
        table_name: "MyTable".to_string(),
    });

    let app = Router::new()
        .route("/scan", get(query_items))
        .with_state(state)
        .layer(TraceLayer::new_for_http());

    axum::Server::bind(&"0.0.0.0:3000"[..])
        .serve(app.into_make_service())
        .await
        .unwrap();
}

Key remediation points:

• Create the DynamoDB client once and wrap it in Arc to share across handlers without per-request allocations.
• Use explicit pagination with a bounded page size and stop accumulation after a safe threshold to prevent unbounded memory growth.
• Avoid storing per-request state in static or global variables; prefer request-local processing and early dropping of large payloads.
• Enable structured logging and monitor RSS to detect gradual growth that may indicate a leak.