test(memtrack): add dlopen/on-demand allocator tests + CI matrix

Author: not-matthias

.github/workflows/ci.yml

@@ -95,7 +95,7 @@ jobs:
         # Each memtrack integration test binary runs its cases serially
         # (eBPF tracker can't overlap with itself in one process), so we
         # shard at the test-binary level to parallelize across jobs.
-        test: [c_tests, cpp_tests, rust_tests, spawn_tests]
+        test: [c_tests, cpp_tests, rust_tests, spawn_tests, dlopen_tests]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:

crates/memtrack/testdata/dlopen/dlopen_alloc.c

@@ -0,0 +1,34 @@
+#include <dlfcn.h>
+#include <stdio.h>
+#include <stddef.h>
+
+/* dlopen an allocator lib and IMMEDIATELY (no sleep) allocate through it. The
+ * immediacy is the race assertion: every allocation must be captured, which
+ * requires the on-demand watcher to have stopped-classified-attached-resumed
+ * before dlopen returns. */
+int main(int argc, char** argv) {
+    if (argc < 2) {
+        fprintf(stderr, "usage: %s <lib>\n", argv[0]);
+        return 1;
+    }
+
+    void* handle = dlopen(argv[1], RTLD_NOW);
+    if (!handle) {
+        fprintf(stderr, "dlopen failed: %s\n", dlerror());
+        return 1;
+    }
+
+    void* (*mi_malloc)(size_t) = (void* (*)(size_t))dlsym(handle, "mi_malloc");
+    void (*mi_free)(void*) = (void (*)(void*))dlsym(handle, "mi_free");
+    if (!mi_malloc || !mi_free) {
+        fprintf(stderr, "dlsym failed\n");
+        return 1;
+    }
+
+    for (int i = 0; i < 100; i++) {
+        void* p = mi_malloc(4242);
+        mi_free(p);
+    }
+
+    return 0;
+}

crates/memtrack/testdata/dlopen/fake_jemalloc.c

@@ -0,0 +1,20 @@
+#include <stddef.h>
+
+/* Self-contained bump allocator; classifies as Jemalloc via the je_malloc
+ * symbol. See fake_mimalloc.c for why it does not call libc malloc. */
+static char arena[1 << 20];
+static size_t off = 0;
+
+void* je_malloc(size_t size) {
+    size_t aligned = (size + 15) & ~((size_t)15);
+    if (off + aligned > sizeof(arena)) {
+        return 0;
+    }
+    void* p = &arena[off];
+    off += aligned;
+    return p;
+}
+
+void je_free(void* p) {
+    (void)p;
+}

crates/memtrack/testdata/dlopen/fake_mimalloc.c

@@ -0,0 +1,21 @@
+#include <stddef.h>
+
+/* Self-contained bump allocator so mi_malloc does not delegate to libc malloc
+ * (which is probed separately and would double-count). Classifies as Mimalloc
+ * via the mi_malloc/mi_free symbols. */
+static char arena[1 << 20];
+static size_t off = 0;
+
+void* mi_malloc(size_t size) {
+    size_t aligned = (size + 15) & ~((size_t)15);
+    if (off + aligned > sizeof(arena)) {
+        return 0;
+    }
+    void* p = &arena[off];
+    off += aligned;
+    return p;
+}
+
+void mi_free(void* p) {
+    (void)p;
+}

crates/memtrack/testdata/dlopen/thread_dlopen.c

@@ -0,0 +1,53 @@
+#include <dlfcn.h>
+#include <pthread.h>
+#include <stddef.h>
+#include <stdio.h>
+
+/* Two threads each dlopen a distinct allocator lib concurrently and allocate
+ * through it. Exercises concurrent stop-the-world: both threads trigger the
+ * watcher, both must be stopped, both libs classified/attached, both resumed. */
+struct job {
+    const char* path;
+    const char* alloc_sym;
+    const char* free_sym;
+    size_t size;
+};
+
+static void* run(void* arg) {
+    struct job* j = (struct job*)arg;
+    void* handle = dlopen(j->path, RTLD_NOW);
+    if (!handle) {
+        return NULL;
+    }
+
+    void* (*alloc)(size_t) = (void* (*)(size_t))dlsym(handle, j->alloc_sym);
+    void (*dealloc)(void*) = (void (*)(void*))dlsym(handle, j->free_sym);
+    if (!alloc || !dealloc) {
+        return NULL;
+    }
+
+    for (int i = 0; i < 100; i++) {
+        void* p = alloc(j->size);
+        dealloc(p);
+    }
+
+    return NULL;
+}
+
+int main(int argc, char** argv) {
+    if (argc < 3) {
+        fprintf(stderr, "usage: %s <mimalloc-lib> <jemalloc-lib>\n", argv[0]);
+        return 1;
+    }
+
+    struct job j1 = {argv[1], "mi_malloc", "mi_free", 4242};
+    struct job j2 = {argv[2], "je_malloc", "je_free", 4243};
+
+    pthread_t t1, t2;
+    pthread_create(&t1, NULL, run, &j1);
+    pthread_create(&t2, NULL, run, &j2);
+    pthread_join(t1, NULL);
+    pthread_join(t2, NULL);
+
+    return 0;
+}

crates/memtrack/tests/cpp_tests.rs

@@ -1,12 +1,30 @@
 #[macro_use]
 mod shared;
 
-use memtrack::AllocatorLib;
 use rstest::rstest;
 use std::path::Path;
 use std::process::Command;
 
+/// A cached cmake configure pins absolute library paths from a previous
+/// environment (e.g. garbage-collected nix store entries). On failure, wipe
+/// the build dir and retry once from a fresh configure.
+///
+/// Builds are serialized: parallel test cases share the build dir, and the
+/// retry path deletes it.
 fn compile_cpp_project(project_dir: &Path, target: &str) -> anyhow::Result<std::path::PathBuf> {
+    static BUILD_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
+    let _guard = BUILD_LOCK.lock().unwrap();
+
+    match build_cpp_target(project_dir, target) {
+        Ok(path) => Ok(path),
+        Err(_) => {
+            std::fs::remove_dir_all(project_dir.join("build"))?;
+            build_cpp_target(project_dir, target)
+        }
+    }
+}
+
+fn build_cpp_target(project_dir: &Path, target: &str) -> anyhow::Result<std::path::PathBuf> {
     let build_exists = project_dir.join("build").exists();
     if !build_exists {
         // Configure with cmake -B build
@@ -57,13 +75,7 @@ fn test_cpp_alloc_tracking(#[case] target: &str) -> Result<(), Box<dyn std::erro
     let project_path = Path::new("testdata/alloc_cpp");
     let binary = compile_cpp_project(project_path, target)?;
 
-    // Try to find a static allocator in the binary, then attach to it as well
-    // This is needed because the CWD is different, which breaks the heuristics.
-    let allocators = AllocatorLib::from_path_static(&binary)
-        .map(|a| vec![a])
-        .unwrap_or_default();
-
-    let (events, thread_handle) = shared::track_binary_with_opts(&binary, &allocators)?;
+    let (events, thread_handle) = shared::track_binary(&binary)?;
     assert_events_with_marker!(target, &events);
 
     thread_handle.join().unwrap();

crates/memtrack/tests/dlopen_tests.rs

@@ -0,0 +1,131 @@
+#[macro_use]
+mod shared;
+
+use runner_shared::artifacts::MemtrackEventKind;
+use std::collections::HashSet;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+use tempfile::TempDir;
+
+fn compile_shared(source: &str, name: &str, dir: &Path) -> PathBuf {
+    let src = dir.join(format!("{name}.c"));
+    std::fs::write(&src, source).expect("write source");
+    let out = dir.join(format!("{name}.so"));
+    let ok = Command::new("gcc")
+        .args(["-shared", "-fPIC", "-o"])
+        .arg(&out)
+        .arg(&src)
+        .status()
+        .expect("run gcc")
+        .success();
+    assert!(ok, "failed to compile {name}.so");
+    out
+}
+
+fn compile_exe(source: &str, name: &str, dir: &Path, libs: &[&str]) -> PathBuf {
+    let src = dir.join(format!("{name}.c"));
+    std::fs::write(&src, source).expect("write source");
+    let out = dir.join(name);
+    let ok = Command::new("gcc")
+        .arg("-o")
+        .arg(&out)
+        .arg(&src)
+        .args(libs)
+        .status()
+        .expect("run gcc")
+        .success();
+    assert!(ok, "failed to compile {name}");
+    out
+}
+
+/// dlopen -> fentry -> SIGSTOP -> resolve -> classify -> attach -> SIGCONT ->
+/// first allocation captured. The dlopen'd allocator's 100 allocations must all
+/// be captured, proving the attach completes before the child's first alloc.
+#[test_with::env(GITHUB_ACTIONS)]
+#[test_log::test]
+fn test_dlopen_allocator() -> Result<(), Box<dyn std::error::Error>> {
+    let dir = TempDir::new()?;
+    let lib = compile_shared(
+        include_str!("../testdata/dlopen/fake_mimalloc.c"),
+        "libfake_mimalloc",
+        dir.path(),
+    );
+    let exe = compile_exe(
+        include_str!("../testdata/dlopen/dlopen_alloc.c"),
+        "dlopen_alloc",
+        dir.path(),
+        &["-ldl"],
+    );
+
+    let mut cmd = Command::new(&exe);
+    cmd.arg(&lib);
+    let (events, thread_handle) = shared::track_command(cmd)?;
+
+    let malloc_addrs: HashSet<u64> = events
+        .iter()
+        .filter_map(|e| match e.kind {
+            MemtrackEventKind::Malloc { size: 4242 } => Some(e.addr),
+            _ => None,
+        })
+        .collect();
+    let malloc_count = events
+        .iter()
+        .filter(|e| matches!(e.kind, MemtrackEventKind::Malloc { size: 4242 }))
+        .count();
+    let free_count = events
+        .iter()
+        .filter(|e| matches!(e.kind, MemtrackEventKind::Free) && malloc_addrs.contains(&e.addr))
+        .count();
+
+    assert_eq!(malloc_count, 100, "expected 100 mi_malloc(4242) events");
+    assert_eq!(
+        free_count, 100,
+        "expected 100 mi_free events for tracked addrs"
+    );
+
+    thread_handle.join().unwrap();
+    Ok(())
+}
+
+/// Two threads dlopen distinct allocator libs concurrently. Both libs' 100
+/// allocations each must be captured, exercising concurrent stop-the-world.
+#[test_with::env(GITHUB_ACTIONS)]
+#[test_log::test]
+fn test_thread_dlopen() -> Result<(), Box<dyn std::error::Error>> {
+    let dir = TempDir::new()?;
+    let mimalloc = compile_shared(
+        include_str!("../testdata/dlopen/fake_mimalloc.c"),
+        "libfake_mimalloc",
+        dir.path(),
+    );
+    let jemalloc = compile_shared(
+        include_str!("../testdata/dlopen/fake_jemalloc.c"),
+        "libfake_jemalloc",
+        dir.path(),
+    );
+    let exe = compile_exe(
+        include_str!("../testdata/dlopen/thread_dlopen.c"),
+        "thread_dlopen",
+        dir.path(),
+        &["-ldl", "-lpthread"],
+    );
+
+    let mut cmd = Command::new(&exe);
+    cmd.arg(&mimalloc).arg(&jemalloc);
+    let (events, thread_handle) = shared::track_command(cmd)?;
+
+    let m4242 = events
+        .iter()
+        .filter(|e| matches!(e.kind, MemtrackEventKind::Malloc { size: 4242 }))
+        .count();
+    let m4243 = events
+        .iter()
+        .filter(|e| matches!(e.kind, MemtrackEventKind::Malloc { size: 4243 }))
+        .count();
+
+    assert_eq!(m4242, 100, "expected 100 mi_malloc(4242) events");
+    assert_eq!(m4243, 100, "expected 100 je_malloc(4243) events");
+
+    thread_handle.join().unwrap();
+    Ok(())
+}

crates/memtrack/tests/rust_tests.rs

@@ -1,7 +1,6 @@
 #[macro_use]
 mod shared;
 
-use memtrack::AllocatorLib;
 use rstest::rstest;
 use std::path::Path;
 
@@ -18,13 +17,8 @@ fn test_rust_alloc_tracking(
     let crate_path = Path::new("testdata/alloc_rust");
     let binary = shared::compile_rust_binary(crate_path, "alloc_rust", features)?;
 
-    // Try to find a static allocator in the binary, then attach to it as well
-    // This is needed because the CWD is different, which breaks the heuristics.
-    let allocators = AllocatorLib::from_path_static(&binary)
-        .map(|a| vec![a])
-        .unwrap_or_default();
-
-    let (events, thread_handle) = shared::track_binary_with_opts(&binary, &allocators)?;
+    // No extra allocators: the watcher must discover the static allocator itself.
+    let (events, thread_handle) = shared::track_binary(&binary)?;
     assert_events_with_marker!(name, &events);
 
     thread_handle.join().unwrap();