feat: support non libc allocators

crates/memtrack/Cargo.toml

@@ -33,6 +33,7 @@ itertools = { workspace = true }
 paste = "1.0"
 libbpf-rs = { version = "0.25.0", features = ["vendored"], optional = true }
 glob = "0.3.3"
+object = { version = "0.36", default-features = false, features = ["read_core", "elf"] }
 
 [build-dependencies]
 libbpf-cargo = { version = "0.25.0", optional = true }

crates/memtrack/src/allocators/dynamic.rs

@@ -0,0 +1,82 @@
+use std::path::PathBuf;
+
+use crate::{AllocatorKind, AllocatorLib};
+
+/// Returns the glob patterns used to find this allocator's shared libraries.
+fn get_allocator_paths(lib: &AllocatorKind) -> &'static [&'static str] {
+    match lib {
+        AllocatorKind::Libc => &[
+            // Debian, Ubuntu: Standard Linux multiarch paths
+            "/lib/*-linux-gnu/libc.so.6",
+            "/usr/lib/*-linux-gnu/libc.so.6",
+            // RHEL, Fedora, CentOS, Arch
+            "/lib*/libc.so.6",
+            "/usr/lib*/libc.so.6",
+            // NixOS: find all glibc versions in the Nix store
+            "/nix/store/*glibc*/lib/libc.so.6",
+        ],
+        AllocatorKind::Jemalloc => &[
+            // Debian, Ubuntu: Standard Linux multiarch paths
+            "/lib/*-linux-gnu/libjemalloc.so*",
+            "/usr/lib/*-linux-gnu/libjemalloc.so*",
+            // RHEL, Fedora, CentOS, Arch
+            "/lib*/libjemalloc.so*",
+            "/usr/lib*/libjemalloc.so*",
+            "/usr/local/lib*/libjemalloc.so*",
+            // NixOS
+            "/nix/store/*jemalloc*/lib/libjemalloc.so*",
+        ],
+        AllocatorKind::Mimalloc => &[
+            // Debian, Ubuntu: Standard Linux multiarch paths
+            "/lib/*-linux-gnu/libmimalloc.so*",
+            "/usr/lib/*-linux-gnu/libmimalloc.so*",
+            // RHEL, Fedora, CentOS, Arch
+            "/lib*/libmimalloc.so*",
+            "/usr/lib*/libmimalloc.so*",
+            "/usr/local/lib*/libmimalloc.so*",
+            // NixOS
+            "/nix/store/*mimalloc*/lib/libmimalloc.so*",
+        ],
+    }
+}
+
+/// Find dynamically linked allocator libraries on the system.
+pub fn find_all() -> anyhow::Result<Vec<AllocatorLib>> {
+    use std::collections::HashSet;
+
+    let mut results = Vec::new();
+    let mut seen_paths: HashSet<PathBuf> = HashSet::new();
+
+    for kind in AllocatorKind::all() {
+        let mut found_any = false;
+
+        for pattern in get_allocator_paths(kind) {
+            let paths = glob::glob(pattern)
+                .ok()
+                .into_iter()
+                .flatten()
+                .filter_map(|p| p.ok())
+                .filter_map(|p| p.canonicalize().ok())
+                .filter(|path| {
+                    std::fs::metadata(path)
+                        .map(|m| m.is_file())
+                        .unwrap_or(false)
+                })
+                .collect::<Vec<_>>();
+
+            for path in paths {
+                if seen_paths.insert(path.clone()) {
+                    results.push(AllocatorLib { kind: *kind, path });
+                    found_any = true;
+                }
+            }
+        }
+
+        // FIXME: Do we still need this?
+        if kind.is_required() && !found_any {
+            anyhow::bail!("Could not find required allocator: {}", kind.name());
+        }
+    }
+
+    Ok(results)
+}

crates/memtrack/src/allocators/mod.rs

@@ -0,0 +1,73 @@
+//! Generic allocator discovery infrastructure.
+//!
+//! This module provides a framework for discovering and attaching to different
+//! memory allocators. It's designed to be easily extensible for adding new allocators.
+
+use std::path::PathBuf;
+
+mod dynamic;
+mod static_linked;
+
+/// Represents the different allocator types we support.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum AllocatorKind {
+    /// Standard C library (glibc, musl, etc.)
+    Libc,
+    /// jemalloc - used by FreeBSD, Firefox, many Rust projects
+    Jemalloc,
+    /// mimalloc - Microsoft's allocator
+    Mimalloc,
+    // Future allocators:
+    // Tcmalloc,
+    // Hoard,
+    // Rpmalloc,
+}
+
+impl AllocatorKind {
+    /// Returns all supported allocator kinds.
+    pub fn all() -> &'static [AllocatorKind] {
+        &[
+            AllocatorKind::Libc,
+            AllocatorKind::Jemalloc,
+            AllocatorKind::Mimalloc,
+        ]
+    }
+
+    /// Returns a human-readable name for the allocator.
+    pub fn name(&self) -> &'static str {
+        match self {
+            AllocatorKind::Libc => "libc",
+            AllocatorKind::Jemalloc => "jemalloc",
+            AllocatorKind::Mimalloc => "mimalloc",
+        }
+    }
+
+    /// Returns true if this allocator is required (must be found).
+    pub fn is_required(&self) -> bool {
+        matches!(self, AllocatorKind::Libc)
+    }
+
+    /// Returns the symbol names used to detect this allocator in binaries.
+    pub fn symbols(&self) -> &'static [&'static str] {
+        match self {
+            AllocatorKind::Libc => &["malloc", "free"],
+            AllocatorKind::Jemalloc => &["_rjem_malloc", "_rjem_free"],
+            AllocatorKind::Mimalloc => &["mi_malloc_aligned", "mi_malloc", "mi_free"],
+        }
+    }
+}
+
+/// Discovered allocator library with its kind and path.
+#[derive(Debug, Clone)]
+pub struct AllocatorLib {
+    pub kind: AllocatorKind,
+    pub path: PathBuf,
+}
+
+impl AllocatorLib {
+    pub fn find_all() -> anyhow::Result<Vec<AllocatorLib>> {
+        let mut allocators = static_linked::find_all()?;
+        allocators.extend(dynamic::find_all()?);
+        Ok(allocators)
+    }
+}

crates/memtrack/src/allocators/static_linked.rs

@@ -0,0 +1,109 @@
+use std::collections::HashSet;
+use std::fs;
+use std::path::{Path, PathBuf};
+
+use crate::allocators::{AllocatorKind, AllocatorLib};
+
+/// Check if a file is an ELF binary by reading its magic bytes.
+fn is_elf(path: &Path) -> bool {
+    let mut file = match fs::File::open(path) {
+        Ok(f) => f,
+        Err(_) => return false,
+    };
+
+    let mut magic = [0u8; 4];
+    use std::io::Read;
+    if file.read_exact(&mut magic).is_err() {
+        return false;
+    }
+
+    // ELF magic: 0x7F 'E' 'L' 'F'
+    magic == [0x7F, b'E', b'L', b'F']
+}
+
+/// Walk upward from current directory to find build directories.
+/// Returns all found build directories in order of preference.
+fn find_build_dirs() -> Vec<PathBuf> {
+    let mut dirs = Vec::new();
+    let Ok(mut current_dir) = std::env::current_dir() else {
+        return dirs;
+    };
+
+    loop {
+        // Check for Cargo/Rust build directory
+        let cargo_analysis = current_dir.join("target").join("codspeed").join("analysis");
+        if cargo_analysis.is_dir() {
+            dirs.push(cargo_analysis);
+        }
+
+        // Check for Bazel build directory
+        let bazel_bin = current_dir.join("bazel-bin");
+        if bazel_bin.is_dir() {
+            dirs.push(bazel_bin);
+        }
+
+        // Check for CMake build directory
+        let cmake_build = current_dir.join("build");
+        if cmake_build.is_dir() {
+            dirs.push(cmake_build);
+        }
+
+        if !current_dir.pop() {
+            break;
+        }
+    }
+
+    dirs
+}
+
+fn find_binaries_in_dir(dir: &Path) -> Vec<PathBuf> {
+    glob::glob(&format!("{}/**/*", dir.display()))
+        .into_iter()
+        .flatten()
+        .filter_map(Result::ok)
+        .filter(|p| p.is_file() && is_elf(p))
+        .collect::<Vec<_>>()
+}
+
+fn find_statically_linked_allocator(path: &Path) -> Option<AllocatorKind> {
+    use object::{Object, ObjectSymbol};
+
+    let data = fs::read(path).ok()?;
+    let file = object::File::parse(&*data).ok()?;
+
+    let symbols: HashSet<_> = file
+        .symbols()
+        .chain(file.dynamic_symbols())
+        .filter(|s| s.is_definition())
+        .filter_map(|s| s.name().ok())
+        .collect();
+
+    // FIXME: We don't support multiple statically linked allocators for now
+
+    AllocatorKind::all()
+        .iter()
+        .find(|kind| kind.symbols().iter().any(|s| symbols.contains(s)))
+        .copied()
+}
+
+pub fn find_all() -> anyhow::Result<Vec<AllocatorLib>> {
+    let build_dirs = find_build_dirs();
+    if build_dirs.is_empty() {
+        return Ok(vec![]);
+    }
+
+    let mut allocators = Vec::new();
+    for build_dir in build_dirs {
+        let bins = find_binaries_in_dir(&build_dir);
+
+        for bin in bins {
+            let Some(kind) = find_statically_linked_allocator(&bin) else {
+                continue;
+            };
+
+            allocators.push(AllocatorLib { kind, path: bin });
+        }
+    }
+
+    Ok(allocators)
+}

crates/memtrack/src/ebpf/c/event.h

@@ -3,13 +3,12 @@
 
 #define EVENT_TYPE_MALLOC 1
 #define EVENT_TYPE_FREE 2
-#define EVENT_TYPE_EXECVE 3
-#define EVENT_TYPE_CALLOC 4
-#define EVENT_TYPE_REALLOC 5
-#define EVENT_TYPE_ALIGNED_ALLOC 6
-#define EVENT_TYPE_MMAP 7
-#define EVENT_TYPE_MUNMAP 8
-#define EVENT_TYPE_BRK 9
+#define EVENT_TYPE_CALLOC 3
+#define EVENT_TYPE_REALLOC 4
+#define EVENT_TYPE_ALIGNED_ALLOC 5
+#define EVENT_TYPE_MMAP 6
+#define EVENT_TYPE_MUNMAP 7
+#define EVENT_TYPE_BRK 8
 
 /* Event structure - shared between BPF and userspace */
 struct event {

crates/memtrack/src/ebpf/c/memtrack.bpf.c

@@ -192,8 +192,8 @@ UPROBE_WITH_ARGS(realloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_REALLO
 /* aligned_alloc: allocates with alignment and size */
 UPROBE_WITH_ARGS(aligned_alloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_ALIGNED_ALLOC)
 
-SEC("tracepoint/syscalls/sys_enter_execve")
-int tracepoint_sys_execve(struct trace_event_raw_sys_enter* ctx) { return submit_event(0, 0, EVENT_TYPE_EXECVE); }
+/* memalign: allocates with alignment and size (legacy interface) */
+UPROBE_WITH_ARGS(memalign, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_ALIGNED_ALLOC)
 
 /* Map to store mmap parameters between entry and return */
 struct mmap_args {

crates/memtrack/src/ebpf/events.rs

@@ -17,7 +17,6 @@ use bindings::*;
 pub enum EventType {
     Malloc = EVENT_TYPE_MALLOC as u8,
     Free = EVENT_TYPE_FREE as u8,
-    Execve = EVENT_TYPE_EXECVE as u8,
     Calloc = EVENT_TYPE_CALLOC as u8,
     Realloc = EVENT_TYPE_REALLOC as u8,
     AlignedAlloc = EVENT_TYPE_ALIGNED_ALLOC as u8,
@@ -31,7 +30,6 @@ impl From<u8> for EventType {
         match val as u32 {
             bindings::EVENT_TYPE_MALLOC => EventType::Malloc,
             bindings::EVENT_TYPE_FREE => EventType::Free,
-            bindings::EVENT_TYPE_EXECVE => EventType::Execve,
             bindings::EVENT_TYPE_CALLOC => EventType::Calloc,
             bindings::EVENT_TYPE_REALLOC => EventType::Realloc,
             bindings::EVENT_TYPE_ALIGNED_ALLOC => EventType::AlignedAlloc,