Persistent Memory Growth / Failure to Deallocate under high Object Store Load (DataFusion + delta-rs)

[eshaanmangal-tf]

Description:

I am experiencing a memory "leak" (unbounded growth) when using snmalloc as the global allocator in a Rust-based query server. The application uses DataFusion and delta-rs to query object storage.

The Issue:

Memory is not being released back to the OS (or reused effectively) at the end of a REST request lifecycle. This leads to a steady climb in RSS until the process is OOM killed.

With jemalloc: Memory is reclaimed/recycled correctly.
With snmalloc: Memory usage climbs linearly and crashes.

Environment:

OS: Linux
snmalloc-rs version: 0.3.8
Relevant Crates: DataFusion, delta-rs (which relies heavily on arrow and FFI).
Runtime: [Tokio 1.48]

Observations:

Interestingly, in the snmalloc trace, the OOM occurs well before the "Limits" (32GB) defined in our monitoring. It seems the allocator is struggling with the specific allocation patterns of DataFusion's execution plan (large buffers for Arrow record batches).

Request:

Are there specific configurations for snmalloc or known issues with the way it interacts with large, short-lived Arrow buffers that might prevent timely deallocation?

Refrence links

delta-io/delta-rs#4241 (comment)

[mjp41]

@eshaanmangal-tf thanks for your report. Would you be able to provide me detailed instructions of how I could repro this? I have had a few issues reported around large allocations and memory growth, but have never been able to replicate.

If there is an easy setup for me to add some tracing to, I would really appreciate help digging into this.

Thanks

[eshaanmangal-tf]

Hi Matthew thanks for the quick reponse. Sharing a code snippet generated by cursor (tbh that was the fastest way)

cargo.toml file

[package]
name = "snmalloc-repro"
version = "0.1.0"
edition = "2024"

[[bin]]
name = "snmalloc_repro"
path = "src/snmalloc_repro/main.rs"

[dependencies]
arrow = "57.3.0"
datafusion = "51.0.0"
deltalake = { version = "0.30.0", features = ["datafusion"] }
tokio = { version = "1.48.0", features = ["full", "rt-multi-thread"] }
url = "2.5.4"

# Default allocator for repro runs
snmalloc-rs = "0.3.8"

# Optional allocator for comparison runs
tikv-jemallocator = { version = "0.6", optional = true }

[features]
default = []
repro-jemalloc = ["dep:tikv-jemallocator"]

commands I used to run code snippet
with snmalloc

cargo run --bin snmalloc_repro -- \
  --mode delta \
  --rows 200000 \
  --payload-bytes 4096 \
  --chunk-rows 256 \
  --iterations 8 \
  --reset-delta \
  --emit-metrics /tmp/snmalloc-delta-metrics.csv \
  --progress-every 25 \
  --cooldown-seconds 30,60,120

with jemalloc

cargo run --bin snmalloc_repro --features repro-jemalloc -- \
  --mode delta \
  --rows 200000 \
  --payload-bytes 4096 \
  --chunk-rows 256 \
  --iterations 8 \
  --reset-delta \
  --emit-metrics /tmp/jemalloc-delta-metrics.csv \
  --progress-every 25 \
  --cooldown-seconds 30,60,120

Mind you here I am not trying to compare snmalloc vs jemalloc, thats not the intent. Intent is to understand deallocations by snmalloc. jemalloc is being used as a ref (or to get ballpark number of how much memory our code can release back to OS)

I can share two CSVs as well if you want to go through metrics before running the code (or in case above code doesnt compile for you)

Code

use std::error::Error;
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Instant;

use datafusion::prelude::{CsvReadOptions, SessionContext};
use deltalake::operations::collect_sendable_stream;
use deltalake::protocol::SaveMode;
use deltalake::DeltaTable;
use tokio::time::{Duration, sleep};
use url::Url;

#[global_allocator]
#[cfg(not(feature = "repro-jemalloc"))]
static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
#[global_allocator]
#[cfg(feature = "repro-jemalloc")]
static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Mode {
    Datafusion,
    Delta,
}

#[derive(Debug, Clone)]
struct Cli {
    mode: Mode,
    csv_path: PathBuf,
    delta_path: PathBuf,
    rows: usize,
    payload_bytes: usize,
    chunk_rows: usize,
    iterations: usize,
    reset_delta: bool,
    emit_metrics: Option<PathBuf>,
    progress_every: usize,
    cooldown_seconds: Vec<u64>,
}

impl Cli {
    fn parse() -> Result<Self, Box<dyn Error>> {
        let mut mode = Mode::Datafusion;
        let mut csv_path = PathBuf::from("/tmp/snmalloc-repro.csv");
        let mut delta_path = PathBuf::from("/tmp/snmalloc-repro-delta");
        let mut rows = 250_000_usize;
        let mut payload_bytes = 2_048_usize;
        let mut chunk_rows = 512_usize;
        let mut iterations = 3_usize;
        let mut reset_delta = false;
        let mut emit_metrics = None;
        let mut progress_every = 100_usize;
        let mut cooldown_seconds = vec![30, 60, 120];

        let mut args = std::env::args().skip(1);
        while let Some(arg) = args.next() {
            match arg.as_str() {
                "--mode" => {
                    let value = args.next().ok_or("--mode requires value")?;
                    mode = match value.as_str() {
                        "datafusion" => Mode::Datafusion,
                        "delta" => Mode::Delta,
                        _ => return Err(format!("invalid --mode {value}, use datafusion|delta").into()),
                    };
                }
                "--csv" => {
                    csv_path = PathBuf::from(args.next().ok_or("--csv requires value")?);
                }
                "--delta" => {
                    delta_path = PathBuf::from(args.next().ok_or("--delta requires value")?);
                }
                "--rows" => {
                    rows = args.next().ok_or("--rows requires value")?.parse()?;
                }
                "--payload-bytes" => {
                    payload_bytes = args.next().ok_or("--payload-bytes requires value")?.parse()?;
                }
                "--chunk-rows" => {
                    chunk_rows = args.next().ok_or("--chunk-rows requires value")?.parse()?;
                }
                "--iterations" => {
                    iterations = args.next().ok_or("--iterations requires value")?.parse()?;
                }
                "--reset-delta" => {
                    reset_delta = true;
                }
                "--emit-metrics" => {
                    emit_metrics = Some(PathBuf::from(
                        args.next().ok_or("--emit-metrics requires value")?,
                    ));
                }
                "--progress-every" => {
                    progress_every = args.next().ok_or("--progress-every requires value")?.parse()?;
                }
                "--cooldown-seconds" => {
                    let raw = args.next().ok_or("--cooldown-seconds requires value")?;
                    cooldown_seconds = parse_cooldown_list(&raw)?;
                }
                "--help" | "-h" => {
                    print_help();
                    std::process::exit(0);
                }
                _ => return Err(format!("unknown arg: {arg}").into()),
            }
        }

        Ok(Self {
            mode,
            csv_path,
            delta_path,
            rows,
            payload_bytes,
            chunk_rows,
            iterations,
            reset_delta,
            emit_metrics,
            progress_every,
            cooldown_seconds,
        })
    }
}

fn print_help() {
    println!(
        "snmalloc_repro

USAGE:
  cargo run --bin snmalloc_repro -- [options]

OPTIONS:
  --mode <datafusion|delta>   Repro mode (default: datafusion)
  --csv <path>                CSV path (default: /tmp/snmalloc-repro.csv)
  --delta <path>              Delta path (default: /tmp/snmalloc-repro-delta)
  --rows <n>                  Rows in generated CSV (default: 250000)
  --payload-bytes <n>         Payload size per row (default: 2048)
  --chunk-rows <n>            Rows per delta file append (default: 512)
  --iterations <n>            Query/scan loops (default: 3)
  --reset-delta               Delete delta path before writing
  --emit-metrics <path>       Write RSS/vmem metrics CSV per iteration
  --progress-every <n>        Print append progress every N files (default: 100)
  --cooldown-seconds <csv>    Post-run memory snapshots, e.g. 30,60,120

EXAMPLES:
  cargo run --bin snmalloc_repro -- --mode datafusion --rows 500000 --payload-bytes 4096
  cargo run --bin snmalloc_repro -- --mode delta --rows 500000 --payload-bytes 4096 --reset-delta
  cargo run --bin snmalloc_repro -- --mode delta --emit-metrics /tmp/snmalloc-metrics.csv
  cargo run --bin snmalloc_repro -- --mode delta --progress-every 50
  cargo run --bin snmalloc_repro --features repro-jemalloc -- --mode delta --emit-metrics /tmp/jemalloc.csv
"
    );
}

fn parse_cooldown_list(raw: &str) -> Result<Vec<u64>, Box<dyn Error>> {
    if raw.trim().is_empty() {
        return Ok(Vec::new());
    }
    let mut values = Vec::new();
    for token in raw.split(',') {
        let token = token.trim();
        if token.is_empty() {
            continue;
        }
        values.push(token.parse::<u64>()?);
    }
    Ok(values)
}

fn ensure_csv(path: &Path, rows: usize, payload_bytes: usize) -> Result<(), Box<dyn Error>> {
    if path.exists() {
        println!("Using existing CSV: {}", path.display());
        return Ok(());
    }
    println!(
        "Generating CSV at {} with rows={} payload_bytes={}",
        path.display(),
        rows,
        payload_bytes
    );
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);
    writer.write_all(b"id,group_key,payload\n")?;

    let payload_seed = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
    let mut payload = String::with_capacity(payload_bytes);
    while payload.len() < payload_bytes {
        payload.push_str(payload_seed);
    }
    payload.truncate(payload_bytes);

    for i in 0..rows {
        // Slightly vary payload to avoid trivially repetitive data.
        let suffix = i % 10_000;
        writeln!(writer, "{i},{},\"{}_{suffix}\"", i % 97, payload)?;
    }
    writer.flush()?;
    Ok(())
}

struct MetricsLogger {
    allocator: &'static str,
    writer: BufWriter<File>,
}

impl MetricsLogger {
    fn new(path: &Path, allocator: &'static str) -> Result<Self, Box<dyn Error>> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let file = File::create(path)?;
        let mut writer = BufWriter::new(file);
        writer.write_all(
            b"unix_ms,allocator,phase,iteration,elapsed_ms,rows,batches,rss_bytes,virtual_bytes\n",
        )?;
        writer.flush()?;
        Ok(Self { allocator, writer })
    }

    fn log(
        &mut self,
        phase: &str,
        iteration: Option<usize>,
        elapsed_ms: Option<u128>,
        rows: usize,
        batches: usize,
    ) -> Result<(), Box<dyn Error>> {
        let now_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)?
            .as_millis();
        let (rss_bytes, vmem_bytes) = capture_process_memory_bytes().unwrap_or((0, 0));
        writeln!(
            self.writer,
            "{now_ms},{},{phase},{},{},{rows},{batches},{rss_bytes},{vmem_bytes}",
            self.allocator,
            iteration.map_or_else(String::new, |v| v.to_string()),
            elapsed_ms.map_or_else(String::new, |v| v.to_string())
        )?;
        self.writer.flush()?;
        Ok(())
    }
}

async fn run_datafusion_mode(cli: &Cli) -> Result<(), Box<dyn Error>> {
    println!("Running datafusion mode");
    let ctx = SessionContext::new();
    ctx.register_csv(
        "fat_csv",
        &cli.csv_path.to_string_lossy(),
        CsvReadOptions::new().has_header(true),
    )
    .await?;

    let mut metrics = if let Some(path) = &cli.emit_metrics {
        Some(MetricsLogger::new(path, current_allocator())?)
    } else {
        None
    };
    if let Some(logger) = &mut metrics {
        logger.log("datafusion_registered_csv", None, None, 0, 0)?;
    }

    for i in 0..cli.iterations {
        let start = Instant::now();
        let df = ctx
            .sql(
                "SELECT group_key, COUNT(*) AS c, MAX(LENGTH(payload)) AS max_payload \
                 FROM fat_csv GROUP BY group_key ORDER BY c DESC",
            )
            .await?;
        let batches = df.collect().await?;
        let rows: usize = batches.iter().map(arrow::record_batch::RecordBatch::num_rows).sum();
        println!(
            "iteration={} batches={} rows={} elapsed_ms={}",
            i,
            batches.len(),
            rows,
            start.elapsed().as_millis()
        );
        if let Some(logger) = &mut metrics {
            logger.log(
                "datafusion_query_collect",
                Some(i),
                Some(start.elapsed().as_millis()),
                rows,
                batches.len(),
            )?;
        }
    }
    record_cooldown_snapshots(&mut metrics, "datafusion", &cli.cooldown_seconds).await?;
    Ok(())
}

async fn create_or_reset_delta_table(cli: &Cli) -> Result<DeltaTable, Box<dyn Error>> {
    if cli.reset_delta && cli.delta_path.exists() {
        println!("Removing existing delta path: {}", cli.delta_path.display());
        fs::remove_dir_all(&cli.delta_path)?;
    }

    let delta_url = path_to_directory_url(&cli.delta_path)?;

    if !cli.delta_path.exists() {
        fs::create_dir_all(&cli.delta_path)?;
        let table = DeltaTable::try_from_url(delta_url.clone())
            .await?
            .create()
            .with_table_name("snmalloc_repro")
            .with_save_mode(SaveMode::Overwrite)
            .with_columns(
                deltalake::kernel::StructType::try_new(vec![
                    deltalake::kernel::StructField::new(
                        "id".to_string(),
                        deltalake::kernel::DataType::Primitive(
                            deltalake::kernel::PrimitiveType::Long,
                        ),
                        false,
                    ),
                    deltalake::kernel::StructField::new(
                        "group_key".to_string(),
                        deltalake::kernel::DataType::Primitive(
                            deltalake::kernel::PrimitiveType::Integer,
                        ),
                        false,
                    ),
                    deltalake::kernel::StructField::new(
                        "payload".to_string(),
                        deltalake::kernel::DataType::Primitive(
                            deltalake::kernel::PrimitiveType::String,
                        ),
                        true,
                    ),
                ])?
                .fields()
                .cloned(),
            )
            .await?;
        return Ok(table);
    }

    Ok(deltalake::open_table(delta_url).await?)
}

async fn run_delta_mode(cli: &Cli) -> Result<(), Box<dyn Error>> {
    println!("Running delta mode");
    let mut metrics = if let Some(path) = &cli.emit_metrics {
        Some(MetricsLogger::new(path, current_allocator())?)
    } else {
        None
    };

    let ctx = SessionContext::new();
    let df = ctx
        .read_csv(
            cli.csv_path.to_string_lossy().as_ref(),
            CsvReadOptions::new().has_header(true),
        )
        .await?;
    let input_batches = df.collect().await?;
    let input_rows: usize = input_batches
        .iter()
        .map(arrow::record_batch::RecordBatch::num_rows)
        .sum();
    println!(
        "Loaded CSV into {} record batches (rows={input_rows})",
        input_batches.len()
    );
    if let Some(logger) = &mut metrics {
        logger.log(
            "delta_loaded_input_csv",
            None,
            None,
            input_rows,
            input_batches.len(),
        )?;
    }

    let mut table = create_or_reset_delta_table(cli).await?;
    let mut files_written = 0_usize;
    let total_files_estimate: usize = input_batches
        .iter()
        .map(|batch| batch.num_rows().div_ceil(cli.chunk_rows))
        .sum();
    println!(
        "Starting delta append: estimated_files={} chunk_rows={} progress_every={}",
        total_files_estimate, cli.chunk_rows, cli.progress_every
    );
    for batch in &input_batches {
        let mut offset = 0_usize;
        while offset < batch.num_rows() {
            let len = (batch.num_rows() - offset).min(cli.chunk_rows);
            let slice = batch.slice(offset, len);
            table = table
                .write(vec![slice])
                .with_save_mode(SaveMode::Append)
                .await?;
            files_written += 1;
            if files_written % cli.progress_every == 0 {
                println!(
                    "append_progress files_written={}/{} ({:.1}%)",
                    files_written,
                    total_files_estimate,
                    (files_written as f64 / total_files_estimate as f64) * 100.0
                );
            }
            offset += len;
        }
    }
    println!(
        "Delta append complete. files_written={} latest_version={:?}",
        files_written,
        table.version()
    );
    if let Some(logger) = &mut metrics {
        logger.log("delta_append_complete", None, None, input_rows, files_written)?;
    }

    for i in 0..cli.iterations {
        let start = Instant::now();
        let (next_table, stream) = table.scan_table().await?;
        let batches = collect_sendable_stream(stream).await?;
        let rows: usize = batches.iter().map(arrow::record_batch::RecordBatch::num_rows).sum();
        println!(
            "scan_iteration={} batches={} rows={} elapsed_ms={}",
            i,
            batches.len(),
            rows,
            start.elapsed().as_millis()
        );
        if let Some(logger) = &mut metrics {
            logger.log(
                "delta_scan_collect",
                Some(i),
                Some(start.elapsed().as_millis()),
                rows,
                batches.len(),
            )?;
        }
        table = next_table;
    }

    record_cooldown_snapshots(&mut metrics, "delta", &cli.cooldown_seconds).await?;
    Ok(())
}

fn path_to_directory_url(path: &Path) -> Result<Url, Box<dyn Error>> {
    let absolute = if path.is_absolute() {
        path.to_path_buf()
    } else {
        std::env::current_dir()?.join(path)
    };
    Url::from_directory_path(absolute)
        .map_err(|_| format!("failed to convert path to directory URL: {}", path.display()).into())
}

fn capture_process_memory_bytes() -> Option<(u64, u64)> {
    #[cfg(target_os = "linux")]
    {
        let status = fs::read_to_string("/proc/self/status").ok()?;
        let rss_kb = parse_status_kb(&status, "VmRSS:")?;
        let vmem_kb = parse_status_kb(&status, "VmSize:")?;
        return Some((rss_kb.saturating_mul(1024), vmem_kb.saturating_mul(1024)));
    }

    #[cfg(not(target_os = "linux"))]
    {
        let pid = std::process::id().to_string();
        let output = Command::new("ps")
            .args(["-o", "rss=", "-o", "vsz=", "-p", &pid])
            .output()
            .ok()?;
        if !output.status.success() {
            return None;
        }
        let text = String::from_utf8(output.stdout).ok()?;
        let mut parts = text.split_whitespace();
        let rss_kb = parts.next()?.parse::<u64>().ok()?;
        let vmem_kb = parts.next()?.parse::<u64>().ok()?;
        Some((rss_kb.saturating_mul(1024), vmem_kb.saturating_mul(1024)))
    }
}

fn current_allocator() -> &'static str {
    if cfg!(feature = "repro-jemalloc") {
        "jemalloc"
    } else {
        "snmalloc"
    }
}

fn align_emit_metrics_path(cli: &mut Cli, allocator: &str) {
    let Some(path) = &mut cli.emit_metrics else {
        return;
    };
    let Some(file_name) = path.file_name().and_then(std::ffi::OsStr::to_str) else {
        return;
    };
    let lower = file_name.to_lowercase();

    let corrected = if allocator == "snmalloc" && lower.contains("jemalloc") {
        Some(file_name.replace("jemalloc", "snmalloc"))
    } else if allocator == "jemalloc" && lower.contains("snmalloc") {
        Some(file_name.replace("snmalloc", "jemalloc"))
    } else {
        None
    };

    if let Some(new_name) = corrected {
        path.set_file_name(&new_name);
        println!(
            "Adjusted --emit-metrics filename for allocator consistency: {}",
            path.display()
        );
    }
}

async fn record_cooldown_snapshots(
    metrics: &mut Option<MetricsLogger>,
    phase_prefix: &str,
    cooldown_seconds: &[u64],
) -> Result<(), Box<dyn Error>> {
    for seconds in cooldown_seconds {
        println!("cooldown_wait seconds={seconds}");
        sleep(Duration::from_secs(*seconds)).await;
        if let Some(logger) = metrics {
            logger.log(
                &format!("{phase_prefix}_cooldown_{seconds}s"),
                None,
                Some((*seconds as u128) * 1000),
                0,
                0,
            )?;
        }
    }
    Ok(())
}

#[cfg(target_os = "linux")]
fn parse_status_kb(status: &str, field: &str) -> Option<u64> {
    for line in status.lines() {
        if line.starts_with(field) {
            let mut it = line.split_whitespace();
            let _ = it.next();
            return it.next()?.parse::<u64>().ok();
        }
    }
    None
}

#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<(), Box<dyn Error>> {
    let mut cli = Cli::parse()?;
    let allocator = current_allocator();
    align_emit_metrics_path(&mut cli, allocator);
    println!("Allocator: {allocator}");
    println!("Config: {cli:#?}");
    ensure_csv(&cli.csv_path, cli.rows, cli.payload_bytes)?;

    match cli.mode {
        Mode::Datafusion => run_datafusion_mode(&cli).await?,
        Mode::Delta => run_delta_mode(&cli).await?,
    }

    println!("Done");
    Ok(())
}