Slowly increasing RSS

I was doing some investigation about concurrent vector search in lancedb causing slowly increasing RSS. Here is a script I have:

#!/usr/bin/env python3

import asyncio
import json
import os
import time

import numpy as np
import psutil

os.environ.setdefault("LANCE_IO_THREADS", "2")
os.environ.setdefault("LANCE_PROCESS_IO_THREADS_LIMIT", "2")
os.environ.setdefault("LANCE_CPU_THREADS", "2")
os.environ.setdefault("AZURE_STORAGE_ACCOUNT_NAME", "aaaaa")
os.environ.setdefault("AZURE_STORAGE_ACCOUNT_KEY", "xxxxx")

import lancedb

URI = "az://examples/fineweb_edu_full_384_dim.lance"
VECTOR_COLUMN = "text_embedding"
VECTOR_DIM = 384
K = 10
CONCURRENCY = 32
REPORT_INTERVAL = 5.0
INDEX_CACHE_SIZE_BYTES = int(1e8)  # 100MB
METADATA_CACHE_SIZE_BYTES = int(1e8)  # 100MB
LOG_PATH = "/tmp/lance_repro.jsonl"
SEARCH_COLUMNS = ["id", "_distance"]


def open_table():
    session = lancedb.Session(
        index_cache_size_bytes=INDEX_CACHE_SIZE_BYTES,
        metadata_cache_size_bytes=METADATA_CACHE_SIZE_BYTES,
    )

    db = lancedb.connect("az://examples", session=session)
    return db.open_table("fineweb_edu_full_384_dim")


async def worker(search_once, stats):
    while True:
        start = time.perf_counter()
        await search_once()
        stats["done"] += 1
        stats["latency_seconds"] += time.perf_counter() - start


async def reporter(stats, log_file, session):
    process = psutil.Process()
    start = time.perf_counter()
    last_done = 0
    last_latency = 0.0
    limit_mb = (INDEX_CACHE_SIZE_BYTES + METADATA_CACHE_SIZE_BYTES) / 1024 / 1024

    while True:
        await asyncio.sleep(REPORT_INTERVAL)
        now = time.perf_counter()
        done = stats["done"]
        latency = stats["latency_seconds"]
        window_done = done - last_done
        window_latency = latency - last_latency
        qps = window_done / REPORT_INTERVAL
        avg_latency_ms = None if window_done == 0 else round(window_latency * 1000 / window_done, 3)
        rss_gb = round(process.memory_info().rss / (1024**3), 3)
        cache_mb = round(session.size_bytes / (1024**2), 1)
        status = "OK" if cache_mb <= limit_mb * 1.1 else "EXCEEDED"

        print(
            f"[{now - start:7.1f}s] done={done:6d} qps={qps:6.2f} avg={avg_latency_ms}ms "
            f"rss={rss_gb:.2f}GB cache={cache_mb:.1f}MB [{status}]",
            flush=True,
        )
        log_file.write(
            json.dumps(
                {
                    "event": "report",
                    "elapsed_seconds": round(now - start, 3),
                    "done": done,
                    "qps": round(qps, 3),
                    "avg_latency_ms": avg_latency_ms,
                    "rss_gb": rss_gb,
                    "cache_mb": cache_mb,
                    "status": status,
                }
            )
            + "\n"
        )
        log_file.flush()

        last_done = done
        last_latency = latency


async def main():
    print("=" * 70)
    print("Memory Growth (cx-style)")
    print("=" * 70)
    print(f"\nConfiguration:")
    print(f"  Dataset: {URI}")
    print(f"  Index cache limit: {INDEX_CACHE_SIZE_BYTES / 1024 / 1024:.0f}MB")
    print(f"  Metadata cache limit: {METADATA_CACHE_SIZE_BYTES / 1024 / 1024:.0f}MB")
    print(f"  Total cache limit: {(INDEX_CACHE_SIZE_BYTES + METADATA_CACHE_SIZE_BYTES) / 1024 / 1024:.0f}MB")
    print(f"  Concurrency: {CONCURRENCY}")
    print()

    session = lancedb.Session(
        index_cache_size_bytes=INDEX_CACHE_SIZE_BYTES,
        metadata_cache_size_bytes=METADATA_CACHE_SIZE_BYTES,
    )
    db = lancedb.connect("az://examples", session=session)

    print("Opening table...")
    table = db.open_table("fineweb_edu_full_384_dim")
    print(f"Table opened: {table.count_rows():,} rows")

    rng = np.random.default_rng()
    log_file = open(LOG_PATH, "w")

    async def search_once():
        await asyncio.to_thread(
            lambda: table.search(
                rng.standard_normal(VECTOR_DIM, dtype=np.float32),
                vector_column_name=VECTOR_COLUMN,
            )
            .select(SEARCH_COLUMNS)
            .limit(K)
            .to_list()
        )

    print(f"\nStarting {CONCURRENCY} concurrent workers...")
    print("=" * 70 + "\n")

    stats = {"done": 0, "latency_seconds": 0.0}
    tasks = [asyncio.create_task(worker(search_once, stats)) for _ in range(CONCURRENCY)]
    tasks.append(asyncio.create_task(reporter(stats, log_file, session)))

    try:
        await asyncio.gather(*tasks)
    except KeyboardInterrupt:
        print("\nInterrupted by user")
    finally:
        log_file.close()


if __name__ == "__main__":
    asyncio.run(main())

This is probably not the best way to do asyncio, but at least this can reproduce the error. Here is an equivalent script using pylance, which I was able to reduce the scope of investigation to pylance becuase it behaves the same:

#!/usr/bin/env python3

import asyncio
import ctypes
import json
import os
import time

import numpy as np
import psutil

os.environ.setdefault("LANCE_IO_THREADS", "2")
os.environ.setdefault("LANCE_PROCESS_IO_THREADS_LIMIT", "2")
os.environ.setdefault("LANCE_CPU_THREADS", "2")
os.environ.setdefault("AZURE_STORAGE_ACCOUNT_NAME", "aaaaaa")
os.environ.setdefault("AZURE_STORAGE_ACCOUNT_KEY", "xxxxxx")

import lance
from lance.lance import _Session as Session

URI = "az://examples/fineweb_edu_full_384_dim.lance"
VECTOR_COLUMN = "text_embedding"
VECTOR_DIM = 384
K = 10
CONCURRENCY = 32
REPORT_INTERVAL = 5.0
INDEX_CACHE_SIZE_BYTES = int(1e8)  # 100MB
METADATA_CACHE_SIZE_BYTES = int(1e8)  # 100MB
LOG_PATH = "/tmp/lance_repro_pylance.jsonl"
SEARCH_COLUMNS = ["id", "_distance"]
MALLOC_TRIM_INTERVAL = 30.0  # Call malloc_trim every 30 seconds

# also add a malloc trim to free up some memory proactively
def get_malloc_trim():
    """Get malloc_trim function from libc if available (Linux only)."""
    try:
        libc = ctypes.CDLL("libc.so.6")
        malloc_trim = libc.malloc_trim
        malloc_trim.argtypes = [ctypes.c_size_t]
        malloc_trim.restype = ctypes.c_int
        return malloc_trim
    except (OSError, AttributeError):
        return None


async def worker(search_once, stats):
    while True:
        start = time.perf_counter()
        await search_once()
        stats["done"] += 1
        stats["latency_seconds"] += time.perf_counter() - start


async def reporter(stats, log_file, session, malloc_trim_fn):
    process = psutil.Process()
    start = time.perf_counter()
    last_done = 0
    last_latency = 0.0
    last_trim_time = start
    limit_mb = (INDEX_CACHE_SIZE_BYTES + METADATA_CACHE_SIZE_BYTES) / 1024 / 1024

    while True:
        await asyncio.sleep(REPORT_INTERVAL)
        now = time.perf_counter()
        done = stats["done"]
        latency = stats["latency_seconds"]
        window_done = done - last_done
        window_latency = latency - last_latency
        qps = window_done / REPORT_INTERVAL
        avg_latency_ms = None if window_done == 0 else round(window_latency * 1000 / window_done, 3)
        rss_before_gb = round(process.memory_info().rss / (1024**3), 3)
        cache_mb = round(session.size_bytes() / (1024**2), 1)
        status = "OK" if cache_mb <= limit_mb * 1.1 else "EXCEEDED"

        trimmed = False
        rss_after_gb = rss_before_gb
        if malloc_trim_fn and (now - last_trim_time) >= MALLOC_TRIM_INTERVAL:
            result = malloc_trim_fn(0)
            trimmed = True
            last_trim_time = now
            rss_after_gb = round(process.memory_info().rss / (1024**3), 3)
            freed_mb = round((rss_before_gb - rss_after_gb) * 1024, 1)
            print(
                f"[{now - start:7.1f}s] done={done:6d} qps={qps:6.2f} avg={avg_latency_ms}ms "
                f"rss={rss_before_gb:.2f}GB -> {rss_after_gb:.2f}GB (freed {freed_mb}MB) cache={cache_mb:.1f}MB [{status}] [TRIMMED]",
                flush=True,
            )
        else:
            print(
                f"[{now - start:7.1f}s] done={done:6d} qps={qps:6.2f} avg={avg_latency_ms}ms "
                f"rss={rss_before_gb:.2f}GB cache={cache_mb:.1f}MB [{status}]",
                flush=True,
            )

        log_file.write(
            json.dumps(
                {
                    "event": "report",
                    "elapsed_seconds": round(now - start, 3),
                    "done": done,
                    "qps": round(qps, 3),
                    "avg_latency_ms": avg_latency_ms,
                    "rss_before_gb": rss_before_gb,
                    "rss_after_gb": rss_after_gb,
                    "cache_mb": cache_mb,
                    "status": status,
                    "trimmed": trimmed,
                }
            )
            + "\n"
        )
        log_file.flush()

        last_done = done
        last_latency = latency


async def main():
    print("=" * 70)
    print("Lance (pylance) Memory Growth")
    print("=" * 70)
    print(f"\nConfiguration:")
    print(f"  Dataset: {URI}")
    print(f"  Index cache limit: {INDEX_CACHE_SIZE_BYTES / 1024 / 1024:.0f}MB")
    print(f"  Metadata cache limit: {METADATA_CACHE_SIZE_BYTES / 1024 / 1024:.0f}MB")
    print(f"  Total cache limit: {(INDEX_CACHE_SIZE_BYTES + METADATA_CACHE_SIZE_BYTES) / 1024 / 1024:.0f}MB")
    print(f"  Concurrency: {CONCURRENCY}")
    print(f"  malloc_trim interval: {MALLOC_TRIM_INTERVAL}s")
    print()

    malloc_trim_fn = get_malloc_trim()
    if malloc_trim_fn:
        print("malloc_trim available - will call every 30 seconds")
    else:
        print("malloc_trim NOT available (not on Linux or libc not found)")
    print()

    # Create session with cache limits
    session = Session(
        index_cache_size_bytes=INDEX_CACHE_SIZE_BYTES,
        metadata_cache_size_bytes=METADATA_CACHE_SIZE_BYTES,
    )

    # Get storage options for Azure
    storage_options = {
        "account_name": os.environ.get("AZURE_STORAGE_ACCOUNT_NAME", ""),
        "account_key": os.environ.get("AZURE_STORAGE_ACCOUNT_KEY", ""),
    }

    print("Opening dataset...")
    ds = lance.dataset(
        URI,
        storage_options=storage_options,
        session=session,
    )
    row_count = ds.count_rows()
    print(f"Dataset opened: {row_count:,} rows")

    rng = np.random.default_rng()
    log_file = open(LOG_PATH, "w")

    async def search_once():
        # Use scanner with nearest parameter for vector search
        # Match lancedb defaults: nprobes=20
        await asyncio.to_thread(
            lambda: ds.scanner(
                columns=SEARCH_COLUMNS,
                nearest={
                    "column": VECTOR_COLUMN,
                    "q": rng.standard_normal(VECTOR_DIM).astype(np.float32),
                    "k": K,
                    "minimum_nprobes": 20,
                    "maximum_nprobes": 20,
                    "use_index": True,
                },
            ).to_table()
        )

    print(f"\nStarting {CONCURRENCY} concurrent workers...")
    print("=" * 70 + "\n")

    stats = {"done": 0, "latency_seconds": 0.0}
    tasks = [asyncio.create_task(worker(search_once, stats)) for _ in range(CONCURRENCY)]
    tasks.append(asyncio.create_task(reporter(stats, log_file, session, malloc_trim_fn)))

    try:
        await asyncio.gather(*tasks)
    except KeyboardInterrupt:
        print("\nInterrupted by user")
    finally:
        log_file.close()


if __name__ == "__main__":
    asyncio.run(main())

Then I was able to reproduce this in lance. But at concurrency 64 this runs very quickly, if I turn concurrency to 2, the growth is more gradual:

//! Equivalent to the Python script repro_pylance.py.
//! This runs concurrent vector searches against a Lance dataset and monitors RSS memory growth.
//!
//! Run with:
//! ```bash
//! export AZURE_STORAGE_ACCOUNT_NAME=aaaaa
//! export AZURE_STORAGE_ACCOUNT_KEY=xxxxxxx
//! ```

#![allow(clippy::print_stdout)]

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use arrow::array::Float32Array;
use clap::Parser;
use futures::TryStreamExt;
use lance::dataset::builder::DatasetBuilder;
use lance::session::Session;
use lance_io::object_store::ObjectStoreRegistry;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use tokio::sync::Semaphore;
use tokio::time::interval;

const URI: &str = "az://examples/fineweb_edu_full_384_dim.lance";
const VECTOR_COLUMN: &str = "text_embedding";
const VECTOR_DIM: usize = 384;
const K: usize = 10;
const CONCURRENCY: usize = 2;
const REPORT_INTERVAL_SECS: f64 = 5.0;
const INDEX_CACHE_SIZE_BYTES: usize = 100_000_000; // 100MB
const METADATA_CACHE_SIZE_BYTES: usize = 100_000_000; // 100MB
const LOG_PATH: &str = "/tmp/lance_repro_rust.jsonl";

#[derive(Parser, Debug)]
#[command(version, about = "Memory leak reproduction benchmark")]
struct Args {
    /// Dataset URI (default: az://examples/fineweb_edu_full_384_dim.lance)
    #[arg(long, default_value = URI)]
    uri: String,

    /// Concurrency level
    #[arg(long, default_value_t = CONCURRENCY)]
    concurrency: usize,

    /// Index cache size in bytes
    #[arg(long, default_value_t = INDEX_CACHE_SIZE_BYTES)]
    index_cache_size: usize,

    /// Metadata cache size in bytes
    #[arg(long, default_value_t = METADATA_CACHE_SIZE_BYTES)]
    metadata_cache_size: usize,
}

/// Get RSS memory in bytes from /proc/self/statm (Linux only)
fn get_rss_bytes() -> Option<u64> {
    #[cfg(target_os = "linux")]
    {
        use std::fs;
        let statm = fs::read_to_string("/proc/self/statm").ok()?;
        let parts: Vec<&str> = statm.split_whitespace().collect();
        if parts.len() >= 2 {
            let rss_pages: u64 = parts[1].parse().ok()?;
            let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) as u64 };
            return Some(rss_pages * page_size);
        }
        None
    }
    #[cfg(not(target_os = "linux"))]
    {
        // For non-Linux, try using sysinfo-like approach or return None
        None
    }
}

/// Call malloc_trim on Linux to release memory back to OS
#[cfg(target_os = "linux")]
fn malloc_trim() -> bool {
    unsafe extern "C" {
        fn malloc_trim(pad: libc::size_t) -> libc::c_int;
    }
    unsafe { malloc_trim(0) != 0 }
}

#[cfg(not(target_os = "linux"))]
fn malloc_trim() -> bool {
    false
}

struct Stats {
    done: AtomicU64,
    latency_nanos: AtomicU64,
}

impl Stats {
    fn new() -> Self {
        Self {
            done: AtomicU64::new(0),
            latency_nanos: AtomicU64::new(0),
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::init();

    let args = Args::parse();

    println!("{}", "=".repeat(70));
    println!("Lance (Rust) Memory Growth Reproduction");
    println!("{}", "=".repeat(70));
    println!("\nConfiguration:");
    println!("  Dataset: {}", args.uri);
    println!(
        "  Index cache limit: {}MB",
        args.index_cache_size / 1024 / 1024
    );
    println!(
        "  Metadata cache limit: {}MB",
        args.metadata_cache_size / 1024 / 1024
    );
    println!(
        "  Total cache limit: {}MB",
        (args.index_cache_size + args.metadata_cache_size) / 1024 / 1024
    );
    println!("  Concurrency: {}", args.concurrency);
    println!();

    // Check if malloc_trim is available
    #[cfg(target_os = "linux")]
    println!("malloc_trim available - will call every 30 seconds");
    #[cfg(not(target_os = "linux"))]
    println!("malloc_trim NOT available (not on Linux)");
    println!();

    // Create session with cache limits
    let session = Arc::new(Session::new(
        args.index_cache_size,
        args.metadata_cache_size,
        Arc::new(ObjectStoreRegistry::default()),
    ));

    // Get storage options from environment
    let mut storage_options = HashMap::new();
    if let Ok(account_name) = std::env::var("AZURE_STORAGE_ACCOUNT_NAME") {
        storage_options.insert("account_name".to_string(), account_name);
    }
    if let Ok(account_key) = std::env::var("AZURE_STORAGE_ACCOUNT_KEY") {
        storage_options.insert("account_key".to_string(), account_key);
    }

    println!("Opening dataset...");
    let dataset = Arc::new(
        DatasetBuilder::from_uri(&args.uri)
            .with_storage_options(storage_options)
            .with_session(session.clone())
            .load()
            .await?,
    );
    let row_count = dataset.count_rows(None).await?;
    println!("Dataset opened: {} rows", row_count);

    let log_file = File::create(LOG_PATH)?;
    let mut log_writer = BufWriter::new(log_file);

    let stats = Arc::new(Stats::new());
    let semaphore = Arc::new(Semaphore::new(args.concurrency));

    println!("\nStarting {} concurrent workers...", args.concurrency);
    println!("{}\n", "=".repeat(70));

    let start = Instant::now();

    // Spawn worker tasks
    let workers: Vec<_> = (0..args.concurrency)
        .map(|_| {
            let dataset = dataset.clone();
            let stats = stats.clone();
            let semaphore = semaphore.clone();

            tokio::spawn(async move {
                let mut rng = StdRng::from_os_rng();
                loop {
                    let _permit = semaphore.acquire().await.unwrap();

                    // Generate random query vector
                    let query: Vec<f32> = (0..VECTOR_DIM)
                        .map(|_| rng.random_range(-1.0..1.0))
                        .collect();
                    let query_array = Float32Array::from(query);

                    let search_start = Instant::now();

                    // Perform vector search
                    let columns: &[&str] = &["id"];
                    let result = dataset
                        .scan()
                        .project(columns)
                        .unwrap()
                        .nearest(VECTOR_COLUMN, &query_array, K)
                        .unwrap()
                        .minimum_nprobes(20)
                        .maximum_nprobes(20)
                        .try_into_stream()
                        .await
                        .unwrap()
                        .try_collect::<Vec<_>>()
                        .await;

                    if let Err(e) = result {
                        eprintln!("Search error: {}", e);
                        continue;
                    }

                    let latency = search_start.elapsed();
                    stats.done.fetch_add(1, Ordering::Relaxed);
                    stats
                        .latency_nanos
                        .fetch_add(latency.as_nanos() as u64, Ordering::Relaxed);
                }
            })
        })
        .collect();

    // Reporter task
    let stats_for_reporter = stats.clone();
    let session_for_reporter = session.clone();
    let reporter = tokio::spawn(async move {
        let mut interval = interval(Duration::from_secs_f64(REPORT_INTERVAL_SECS));
        let mut last_done: u64 = 0;
        let mut last_latency: u64 = 0;
        let mut last_trim_time = Instant::now();
        let limit_mb =
            (INDEX_CACHE_SIZE_BYTES + METADATA_CACHE_SIZE_BYTES) as f64 / 1024.0 / 1024.0;

        loop {
            interval.tick().await;

            let now = Instant::now();
            let elapsed = start.elapsed().as_secs_f64();
            let done = stats_for_reporter.done.load(Ordering::Relaxed);
            let latency = stats_for_reporter.latency_nanos.load(Ordering::Relaxed);

            let window_done = done - last_done;
            let window_latency = latency - last_latency;
            let qps = window_done as f64 / REPORT_INTERVAL_SECS;
            let avg_latency_ms = if window_done > 0 {
                Some((window_latency as f64 / window_done as f64) / 1_000_000.0)
            } else {
                None
            };

            let rss_before = get_rss_bytes().unwrap_or(0);
            let rss_before_gb = rss_before as f64 / (1024.0 * 1024.0 * 1024.0);
            let cache_bytes = session_for_reporter.size_bytes();
            let cache_mb = cache_bytes as f64 / (1024.0 * 1024.0);
            let status = if cache_mb <= limit_mb * 1.1 {
                "OK"
            } else {
                "EXCEEDED"
            };

            // Check if we should call malloc_trim (every 30 seconds)
            let trimmed = if now.duration_since(last_trim_time).as_secs() >= 30 {
                let result = malloc_trim();
                last_trim_time = now;
                result
            } else {
                false
            };

            let rss_after = get_rss_bytes().unwrap_or(0);
            let rss_after_gb = rss_after as f64 / (1024.0 * 1024.0 * 1024.0);

            if trimmed {
                let freed_mb = (rss_before as f64 - rss_after as f64) / (1024.0 * 1024.0);
                println!(
                    "[{:7.1}s] done={:6} qps={:6.2} avg={:.3}ms rss={:.2}GB -> {:.2}GB (freed {:.1}MB) cache={:.1}MB [{}] [TRIMMED]",
                    elapsed,
                    done,
                    qps,
                    avg_latency_ms.unwrap_or(0.0),
                    rss_before_gb,
                    rss_after_gb,
                    freed_mb,
                    cache_mb,
                    status
                );
            } else {
                println!(
                    "[{:7.1}s] done={:6} qps={:6.2} avg={:.3}ms rss={:.2}GB cache={:.1}MB [{}]",
                    elapsed,
                    done,
                    qps,
                    avg_latency_ms.unwrap_or(0.0),
                    rss_before_gb,
                    cache_mb,
                    status
                );
            }

            // Write to log file
            let log_entry = serde_json::json!({
                "event": "report",
                "elapsed_seconds": elapsed,
                "done": done,
                "qps": qps,
                "avg_latency_ms": avg_latency_ms,
                "rss_before_gb": rss_before_gb,
                "rss_after_gb": rss_after_gb,
                "cache_mb": cache_mb,
                "status": status,
                "trimmed": trimmed,
            });
            writeln!(log_writer, "{}", log_entry).ok();
            log_writer.flush().ok();

            last_done = done;
            last_latency = latency;
        }
    });

    // Wait for Ctrl+C
    tokio::select! {
        _ = tokio::signal::ctrl_c() => {
            println!("\nInterrupted by user");
        }
        _ = async {
            for worker in workers {
                worker.await.ok();
            }
        } => {}
    }

    reporter.abort();

    Ok(())
}