Profiling Heap Allocations in Production Rust Services with Jemalloc and Grafana Pyroscope — code snippets

snippet-1.txt

// Cargo.toml configuration:
// [dependencies]
// tikv-jemallocator = { version = "0.6", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] }
// tikv-jemalloc-ctl = "0.6"
// tokio = { version = "1.0", features = ["full"] }

use std::alloc::System;

#[global_allocator]
static GLOBAL: tikv-jemallocator::Jemalloc = tikv-jemallocator::Jemalloc;

fn main() {
    // Application initialization logic
    println!("Jemalloc-enabled production service started.");
}

snippet-2.txt

use dashmap::DashMap;
use std::sync::Arc;
use tokio::time::{sleep, Duration};

lazy_static::lazy_static! {
    static ref METADATA_CACHE: Arc<DashMap<String, Vec<u8>>> = Arc::new(DashMap::new());
}

async fn process_request(request_id: String) {
    // Simulate parsing a large payload (1 MB)
    let heavy_payload = vec![0u8; 1_048_576]; 
    
    // Cache the payload for processing
    METADATA_CACHE.insert(request_id.clone(), heavy_payload);
    
    // Simulate business logic processing
    sleep(Duration::from_millis(50)).await;

    // BUG: Under error conditions, we return early and skip cache eviction
    if request_id.contains("leak") {
        return; // Leaks the 1MB payload in the global cache
    }

    // Normal path eviction
    METADATA_CACHE.remove(&request_id);
}

#[tokio::main]
async fn main() {
    let mut counter = 0;
    loop {
        counter += 1;
        let id = if counter % 10 == 0 {
            format!("req-leak-{}", counter)
        } else {
            format!("req-normal-{}", counter)
        };
        tokio::spawn(process_request(id));
        sleep(Duration::from_millis(10)).await;
    }
}

snippet-3.txt

use axum::{routing::post, Router, response::IntoResponse, http::StatusCode};
use std::ffi::CString;

async fn trigger_heap_dump() -> impl IntoResponse {
    // Convert target file path to a C-compatible string
    let file_path = CString::new("/tmp/jemalloc_profile.out").unwrap();
    
    unsafe {
        // Trigger raw Jemalloc profile dump
        match tikv_jemalloc_ctl::raw::write(
            b"prof.dump\0", 
            file_path.as_ptr()
        ) {
            Ok(_) => (StatusCode::OK, "Heap profile written to /tmp/jemalloc_profile.out"),
            Err(e) => {
                eprintln!("Failed to write heap dump: {:?}", e);
                (StatusCode::INTERNAL_SERVER_ERROR, "Failed to write heap dump")
            }
        }
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/debug/pprof/heap", post(trigger_heap_dump));
    let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

snippet-4.sh

# Run the binary with profiling enabled but inactive, or active by default
export MALLOC_CONF="prof:true,prof_active:true,lg_prof_sample:19"
./target/release/my_rust_service

snippet-5.txt

use pyroscope::{PyroscopeAgent, pyroscope::PyroscopeAgentBuilder};
use pyroscope_jemalloc::JemallocBackend;
use std::env;

fn initialize_continuous_profiling() -> Result<PyroscopeAgent, Box<dyn std::error::Error>> {
    let pyroscope_url = env::var("PYROSCOPE_SERVER_URL")
        .unwrap_or_else(|_| "http://localhost:4040".to_string());

    // Build the Pyroscope agent with the Jemalloc profiling backend
    let agent = PyroscopeAgentBuilder::new(&pyroscope_url, "rust-payment-service")
        .backend(JemallocBackend::default())
        .tags([("env", "production"), ("region", "us-west-2")].to_vec())
        .build()?;

    // Start the agent background loop
    let active_agent = agent.start()?;
    Ok(active_agent)
}

fn main() {
    // Start continuous profiling
    let _agent = initialize_continuous_profiling().expect("Failed to initialize Pyroscope agent");

    // Start application loop...
}

snippet-6.yaml

version: '3.8'

services:
  pyroscope:
    image: grafana/pyroscope:latest
    ports:
      - "4040:4040"
    volumes:
      - pyroscope-data:/var/lib/pyroscope
    command: ["-config.file=/etc/pyroscope/pyroscope.yaml"]

  rust-service:
    build:
      context: .
      dockerfile: Dockerfile
    environment:
      - PYROSCOPE_SERVER_URL=http://pyroscope:4040
      - MALLOC_CONF=prof:true,prof_active:true,lg_prof_sample:19
    ports:
      - "8080:8080"
    depends_on:
      - pyroscope

volumes:
  pyroscope-data:

snippet-7.sh

# Step 1: Capture baseline heap dump
curl -X POST http://localhost:8080/debug/pprof/heap
mv /tmp/jemalloc_profile.out ./baseline.heap

# Wait for memory leak to accumulate...

# Step 2: Capture leak heap dump
curl -X POST http://localhost:8080/debug/pprof/heap
mv /tmp/jemalloc_profile.out ./leaking.heap

# Step 3: Compare both dumps using jeprof to isolate the delta
jeprof --show_bytes --pdf \
  ./target/release/my_rust_service \
  --base=./baseline.heap \
  ./leaking.heap > diff.pdf