feat(memtrack): support dynamically and statically linked allocators

Author: not-matthias

Cargo.lock

@@ -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/Cargo.toml

@@ -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/dynamic.rs

@@ -0,0 +1,64 @@
+//! 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)
+    }
+}
+
+/// 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/mod.rs

@@ -0,0 +1,77 @@
+use anyhow::Context;
+use std::collections::HashSet;
+use std::fs;
+use std::path::{Path, PathBuf};
+
+use crate::allocators::{AllocatorKind, AllocatorLib};
+
+const JEMALLOC_SYMBOLS: &[&str] = &["_rjem_malloc", "_rjem_free"];
+const MIMALLOC_SYMBOLS: &[&str] = &["mi_malloc_aligned", "mi_malloc", "mi_free"];
+
+/// Walk upward from current directory to find `target/codspeed/analysis`.
+fn find_analysis_dir() -> Option<PathBuf> {
+    let mut dir = std::env::current_dir().ok()?;
+
+    loop {
+        let candidate = dir.join("target").join("codspeed").join("analysis");
+        if candidate.is_dir() {
+            return Some(candidate);
+        }
+        if !dir.pop() {
+            return None;
+        }
+    }
+}
+
+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())
+        .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 dont support multiple statically linked allocators for now
+
+    if JEMALLOC_SYMBOLS.iter().any(|s| symbols.contains(s)) {
+        Some(AllocatorKind::Jemalloc)
+    } else if MIMALLOC_SYMBOLS.iter().any(|s| symbols.contains(s)) {
+        Some(AllocatorKind::Mimalloc)
+    } else {
+        None
+    }
+}
+
+pub fn find_all() -> anyhow::Result<Vec<AllocatorLib>> {
+    // 1. Find all the binaries built by CodSpeed
+    let bins =
+        find_binaries_in_dir(&find_analysis_dir().context("Failed to find analysis directory")?);
+
+    let mut allocators = Vec::new();
+    for bin in bins {
+        let Some(kind) = find_statically_linked_allocator(&bin) else {
+            continue;
+        };
+
+        allocators.push(AllocatorLib {
+            kind,
+            path: bin.to_path_buf(),
+        });
+    }
+
+    Ok(allocators)
+}