diff --git a/Cargo.lock b/Cargo.lock index 28bdcbf6..fbb557bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1658,6 +1658,7 @@ dependencies = [ "libbpf-rs", "libc", "log", + "object 0.36.7", "paste", "rstest 0.21.0", "runner-shared", diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index 2d739890..1e713830 100644 --- a/crates/memtrack/Cargo.toml +++ b/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 } diff --git a/crates/memtrack/src/allocators/dynamic.rs b/crates/memtrack/src/allocators/dynamic.rs new file mode 100644 index 00000000..c35e427c --- /dev/null +++ b/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> { + use std::collections::HashSet; + + let mut results = Vec::new(); + let mut seen_paths: HashSet = 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::>(); + + 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) +} diff --git a/crates/memtrack/src/allocators/mod.rs b/crates/memtrack/src/allocators/mod.rs new file mode 100644 index 00000000..70ad34ec --- /dev/null +++ b/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> { + let mut allocators = static_linked::find_all()?; + allocators.extend(dynamic::find_all()?); + Ok(allocators) + } +} diff --git a/crates/memtrack/src/allocators/static_linked.rs b/crates/memtrack/src/allocators/static_linked.rs new file mode 100644 index 00000000..2a0f9bab --- /dev/null +++ b/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 { + 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 { + glob::glob(&format!("{}/**/*", dir.display())) + .into_iter() + .flatten() + .filter_map(Result::ok) + .filter(|p| p.is_file() && is_elf(p)) + .collect::>() +} + +fn find_statically_linked_allocator(path: &Path) -> Option { + 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> { + 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) +} diff --git a/crates/memtrack/src/ebpf/c/event.h b/crates/memtrack/src/ebpf/c/event.h index 0162078d..5204df59 100644 --- a/crates/memtrack/src/ebpf/c/event.h +++ b/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 { diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 63c5be5f..b6927f0d 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/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 { diff --git a/crates/memtrack/src/ebpf/events.rs b/crates/memtrack/src/ebpf/events.rs index f8edbc8f..204d41c0 100644 --- a/crates/memtrack/src/ebpf/events.rs +++ b/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 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, diff --git a/crates/memtrack/src/ebpf/memtrack.rs b/crates/memtrack/src/ebpf/memtrack.rs index c9ac47da..d444eea2 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/crates/memtrack/src/ebpf/memtrack.rs @@ -1,14 +1,13 @@ -use anyhow::Context; -use anyhow::Result; +use crate::prelude::*; use libbpf_rs::Link; use libbpf_rs::skel::OpenSkel; use libbpf_rs::skel::SkelBuilder; use libbpf_rs::{MapCore, UprobeOpts}; -use log::warn; use paste::paste; use std::mem::MaybeUninit; use std::path::Path; +use crate::allocators::AllocatorKind; use crate::ebpf::poller::RingBufferPoller; pub mod memtrack_skel { @@ -16,100 +15,139 @@ pub mod memtrack_skel { } pub use memtrack_skel::*; -/// Macro to attach a function with both entry and return probes +/// Resolve symbol offset from .symtab to ensure that libbpf can find it. Otherwise +/// it will print a warning at runtime. +fn ensure_symbol_exists(lib_path: &Path, symbol_name: &str) -> Result<()> { + use object::{Object, ObjectSymbol}; + + let data = std::fs::read(lib_path)?; + let file = object::File::parse(&*data)?; + + // Check both regular and dynamic symbols + for symbol in file.symbols().chain(file.dynamic_symbols()) { + if !symbol.is_definition() { + continue; + } + + let Ok(name) = symbol.name() else { + continue; + }; + + if name == symbol_name { + let addr = symbol.address(); + if addr != 0 { + return Ok(()); + } + } + } + + bail!("Symbol {symbol_name} not found in {}", lib_path.display()) +} + +/// Macro to attach a function with both entry and return probes. +/// Also generates a `try_attach_*` variant that logs errors instead of returning them. +/// +/// Uses offset-based attachment by resolving symbols from .symtab. +/// Fails if the symbol is not found. macro_rules! attach_uprobe_uretprobe { - ($name:ident, $prog_entry:ident, $prog_return:ident, $func_str:expr) => { - fn $name(&mut self, libc_path: &Path) -> Result<()> { + ($name:ident, $prog_entry:ident, $prog_return:ident) => { + fn $name(&mut self, lib_path: &Path, symbol: &str) -> Result<()> { + ensure_symbol_exists(lib_path, symbol)?; + + // Attach entry probe at function entry via func_name let link = self .skel .progs .$prog_entry .attach_uprobe_with_opts( -1, - libc_path, + lib_path, 0, UprobeOpts { - func_name: Some($func_str.to_string()), retprobe: false, + func_name: Some(symbol.to_owned()), ..Default::default() }, ) .context(format!( "Failed to attach {} uprobe in {}", - $func_str, - libc_path.display() + symbol, + lib_path.display() ))?; self.probes.push(link); + // Attach return probe at function entry via func_name let link = self .skel .progs .$prog_return .attach_uprobe_with_opts( -1, - libc_path, + lib_path, 0, UprobeOpts { - func_name: Some($func_str.to_string()), retprobe: true, + func_name: Some(symbol.to_owned()), ..Default::default() }, ) .context(format!( "Failed to attach {} uretprobe in {}", - $func_str, - libc_path.display() + symbol, + lib_path.display() ))?; self.probes.push(link); Ok(()) } - }; - ($name:ident) => { + paste! { - attach_uprobe_uretprobe!( - [], - [], - [], - stringify!($name) - ); + fn [](&mut self, lib_path: &Path, symbol: &str) { + let result = self.$name(lib_path, symbol); + log::trace!("{} uprobe attach result: {:?}", symbol, result); + } } }; } +/// Macro to attach a function with only an entry probe (no return probe). +/// Also generates a `try_attach_*` variant that logs errors instead of returning them. +/// +/// Uses offset-based attachment by resolving symbols from .symtab. +/// Fails if the symbol is not found. macro_rules! attach_uprobe { - ($name:ident, $prog:ident, $func_str:expr) => { - fn $name(&mut self, libc_path: &Path) -> Result<()> { + ($name:ident, $prog:ident) => { + fn $name(&mut self, lib_path: &Path, symbol: &str) -> Result<()> { + ensure_symbol_exists(lib_path, symbol)?; + let link = self .skel .progs .$prog .attach_uprobe_with_opts( -1, - libc_path, + lib_path, 0, UprobeOpts { - func_name: Some($func_str.to_string()), retprobe: false, + func_name: Some(symbol.to_owned()), ..Default::default() }, ) .context(format!( "Failed to attach {} uprobe in {}", - $func_str, - libc_path.display() + symbol, + lib_path.display() ))?; self.probes.push(link); Ok(()) } - }; - ($name:ident) => { + paste! { - attach_uprobe!( - [], - [], - stringify!($name) - ); + fn [](&mut self, lib_path: &Path, symbol: &str) { + let result = self.$name(lib_path, symbol); + log::trace!("{} uprobe attach result: {:?}", symbol, result); + } } }; } @@ -206,21 +244,111 @@ impl MemtrackBpf { Ok(()) } - attach_uprobe_uretprobe!(malloc); - attach_uprobe_uretprobe!(calloc); - attach_uprobe_uretprobe!(realloc); - attach_uprobe_uretprobe!(aligned_alloc); - attach_uprobe!(free); - - pub fn attach_probes(&mut self, libc_path: &Path) -> Result<()> { - self.attach_malloc(libc_path)?; - self.attach_free(libc_path)?; - self.attach_calloc(libc_path)?; - self.attach_realloc(libc_path)?; - self.attach_aligned_alloc(libc_path)?; + // ========================================================================= + // Allocation probe functions (symbol passed at call time) + // ========================================================================= + attach_uprobe_uretprobe!(attach_malloc, uprobe_malloc, uretprobe_malloc); + attach_uprobe_uretprobe!(attach_calloc, uprobe_calloc, uretprobe_calloc); + attach_uprobe_uretprobe!(attach_realloc, uprobe_realloc, uretprobe_realloc); + attach_uprobe_uretprobe!( + attach_aligned_alloc, + uprobe_aligned_alloc, + uretprobe_aligned_alloc + ); + attach_uprobe_uretprobe!(attach_memalign, uprobe_memalign, uretprobe_memalign); + attach_uprobe!(attach_free, uprobe_free); + + // ========================================================================= + // Attach methods grouped by allocator + // ========================================================================= + + /// Attach standard library allocation probes (libc-style: malloc, free, calloc, etc.) + /// This works for libc and allocators that export standard symbol names. + /// For non-libc allocators, standard names are optional - just try them silently. + pub fn attach_libc_probes(&mut self, lib_path: &Path) -> Result<()> { + self.try_attach_malloc(lib_path, "malloc"); + self.try_attach_calloc(lib_path, "calloc"); + self.try_attach_realloc(lib_path, "realloc"); + self.try_attach_free(lib_path, "free"); + self.try_attach_aligned_alloc(lib_path, "aligned_alloc"); + self.try_attach_memalign(lib_path, "memalign"); + Ok(()) + } + + /// Attach probes for a specific allocator kind. + /// This attaches both standard probes (if the allocator exports them) and + /// allocator-specific prefixed probes. + pub fn attach_allocator_probes(&mut self, kind: AllocatorKind, lib_path: &Path) -> Result<()> { + debug!( + "Attaching {} probes to: {}", + kind.name(), + lib_path.display() + ); + + match kind { + AllocatorKind::Libc => { + // Libc only has standard probes, and they must succeed + self.attach_libc_probes(lib_path) + } + AllocatorKind::Jemalloc => { + // Try standard names (jemalloc may export these as drop-in replacements) + let _ = self.attach_libc_probes(lib_path); + self.attach_jemalloc_probes(lib_path) + } + AllocatorKind::Mimalloc => { + // Try standard names (mimalloc may export these as drop-in replacements) + let _ = self.attach_libc_probes(lib_path); + self.attach_mimalloc_probes(lib_path) + } + } + } + + /// Attach jemalloc-specific probes (prefixed and extended API). + fn attach_jemalloc_probes(&mut self, lib_path: &Path) -> Result<()> { + // The following functions are used in Rust when setting a global allocator: + // - rust_alloc: _rjem_malloc and _rjem_mallocx + // - rust_alloc_zeroed: _rjem_mallocx / _rjem_calloc + // - rust_dealloc: _rjem_sdallocx + // - rust_realloc: _rjem_realloc / _rjem_rallocx + + // Prefixed standard API + self.try_attach_malloc(lib_path, "_rjem_malloc"); + self.try_attach_malloc(lib_path, "_rjem_mallocx"); // Also used for `calloc` + self.try_attach_calloc(lib_path, "_rjem_calloc"); + self.try_attach_realloc(lib_path, "_rjem_realloc"); + self.try_attach_realloc(lib_path, "_rjem_rallocx"); + self.try_attach_aligned_alloc(lib_path, "_rjem_aligned_alloc"); + self.try_attach_memalign(lib_path, "_rjem_memalign"); + self.try_attach_free(lib_path, "_rjem_free"); + self.try_attach_free(lib_path, "_rjem_sdallocx"); + Ok(()) } + /// Attach mimalloc-specific probes (mi_* API). + fn attach_mimalloc_probes(&mut self, lib_path: &Path) -> Result<()> { + // The following functions are used in Rust when setting a global allocator: + // - mi_malloc_aligned + // - mi_free + // - mi_realloc_aligned + // - mi_zalloc_aligned + + // Core API + self.try_attach_malloc(lib_path, "mi_malloc"); + self.try_attach_malloc(lib_path, "mi_malloc_aligned"); + self.try_attach_calloc(lib_path, "mi_calloc"); + self.try_attach_realloc(lib_path, "mi_realloc"); + self.try_attach_aligned_alloc(lib_path, "mi_aligned_alloc"); + self.try_attach_memalign(lib_path, "mi_memalign"); + self.try_attach_free(lib_path, "mi_free"); + + // Zero-initialized and aligned variants + self.try_attach_calloc(lib_path, "mi_zalloc"); + self.try_attach_calloc(lib_path, "mi_zalloc_aligned"); + self.try_attach_realloc(lib_path, "mi_realloc_aligned"); + + Ok(()) + } attach_tracepoint!(sched_fork); attach_tracepoint!(sys_execve); diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 5c864918..683135a4 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -1,6 +1,8 @@ -use crate::ebpf::{Event, MemtrackBpf}; -use anyhow::Result; -use log::debug; +use crate::prelude::*; +use crate::{ + AllocatorLib, + ebpf::{Event, MemtrackBpf}, +}; use std::sync::mpsc::{self, Receiver}; pub struct Tracker { @@ -19,16 +21,16 @@ impl Tracker { // Bump memlock limits Self::bump_memlock_rlimit()?; - // Find and attach to all libc instances - let libc_paths = crate::libc::find_libc_paths()?; - debug!("Found {} libc instance(s)", libc_paths.len()); - let mut bpf = MemtrackBpf::new()?; bpf.attach_tracepoints()?; - for libc_path in &libc_paths { - debug!("Attaching uprobes to: {}", libc_path.display()); - bpf.attach_probes(libc_path)?; + // Find and attach to all allocators + let allocators = AllocatorLib::find_all()?; + debug!("Found {} allocator instance(s)", allocators.len()); + + for allocator in &allocators { + debug!("Attaching uprobes to: {}", allocator.path.display()); + bpf.attach_allocator_probes(allocator.kind, &allocator.path)?; } Ok(Self { bpf }) diff --git a/crates/memtrack/src/ipc.rs b/crates/memtrack/src/ipc.rs index 8c4c3b94..2b0d80d5 100644 --- a/crates/memtrack/src/ipc.rs +++ b/crates/memtrack/src/ipc.rs @@ -1,4 +1,4 @@ -use anyhow::{Context, Result}; +use crate::prelude::*; use ipc_channel::ipc::{self, IpcOneShotServer, IpcSender}; use serde::{Deserialize, Serialize}; @@ -86,8 +86,6 @@ impl MemtrackIpcClient { /// Handle incoming IPC messages in memtrack #[cfg(feature = "ebpf")] pub fn handle_ipc_message(msg: IpcMessage, tracker: &std::sync::Arc>) { - use log::debug; - let response = match msg.command { IpcCommand::Enable => match tracker.lock() { Ok(mut t) => match t.enable() { diff --git a/crates/memtrack/src/lib.rs b/crates/memtrack/src/lib.rs index b62ed344..dc32092a 100644 --- a/crates/memtrack/src/lib.rs +++ b/crates/memtrack/src/lib.rs @@ -1,7 +1,10 @@ +mod allocators; #[cfg(feature = "ebpf")] mod ebpf; mod ipc; -mod libc; +pub mod prelude; + +pub use allocators::{AllocatorKind, AllocatorLib}; pub use ipc::{ IpcCommand as MemtrackIpcCommand, IpcMessage as MemtrackIpcMessage, IpcResponse as MemtrackIpcResponse, MemtrackIpcClient, MemtrackIpcServer, diff --git a/crates/memtrack/src/libc.rs b/crates/memtrack/src/libc.rs deleted file mode 100644 index bc0a7feb..00000000 --- a/crates/memtrack/src/libc.rs +++ /dev/null @@ -1,36 +0,0 @@ -#[cfg(feature = "ebpf")] -pub fn find_libc_paths() -> anyhow::Result> { - use itertools::Itertools; - - let patterns = [ - // 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", - ]; - - let existing_paths = patterns - .iter() - .flat_map(|pattern| glob::glob(pattern).ok()) - .flatten() - .filter_map(|p| p.ok()) - .filter_map(|p| p.canonicalize().ok()) - .filter(|path| { - let Ok(metadata) = std::fs::metadata(path) else { - return false; - }; - metadata.is_file() - }) - .dedup() - .collect::>(); - - if existing_paths.is_empty() { - anyhow::bail!("Could not find libc.so.6"); - } - - Ok(existing_paths) -} diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index fa7a64ce..02019900 100644 --- a/crates/memtrack/src/main.rs +++ b/crates/memtrack/src/main.rs @@ -1,9 +1,9 @@ -use anyhow::{Context, Result, anyhow}; use clap::Parser; -use ipc_channel::ipc::{self}; -use log::{debug, info}; +use ipc_channel::ipc; +use memtrack::prelude::*; use memtrack::{MemtrackIpcMessage, Tracker, handle_ipc_message}; use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, MemtrackEvent, MemtrackWriter}; +use std::os::unix::process::CommandExt; use std::path::{Path, PathBuf}; use std::process::Command; use std::sync::atomic::{AtomicBool, Ordering}; @@ -37,6 +37,14 @@ enum Commands { }, } +/// Get the original user's UID and GID when running under sudo. +/// Returns None if not running under sudo or if the environment variables are not set. +fn get_user_uid_gid() -> Option<(u32, u32)> { + let uid = std::env::var("SUDO_UID").ok()?.parse().ok()?; + let gid = std::env::var("SUDO_GID").ok()?.parse().ok()?; + Some((uid, gid)) +} + fn main() -> Result<()> { env_logger::builder() .parse_env(env_logger::Env::new().filter_or("CODSPEED_LOG", "info")) @@ -66,16 +74,25 @@ fn track_command( ipc_server_name: Option, out_dir: &Path, ) -> anyhow::Result { - let tracker = Tracker::new()?; - - let tracker_arc = Arc::new(Mutex::new(tracker)); - let ipc_handle = if let Some(server_name) = ipc_server_name { + // First, establish IPC connection if needed to avoid timeouts on the runner because + // creating the Tracker instance takes some time. + let ipc_channel = if let Some(server_name) = ipc_server_name { debug!("Connecting to IPC server: {server_name}"); let (tx, rx) = ipc::channel::()?; let sender = ipc::IpcSender::connect(server_name)?; sender.send(tx)?; + Some(rx) + } else { + None + }; + + let tracker = Tracker::new()?; + let tracker_arc = Arc::new(Mutex::new(tracker)); + + // Spawn IPC handler thread with the now-available tracker + let ipc_handle = if let Some(rx) = ipc_channel { let tracker_clone = tracker_arc.clone(); Some(thread::spawn(move || { while let Ok(msg) = rx.recv() { @@ -87,9 +104,18 @@ fn track_command( }; // Start the target command using bash to handle shell syntax - let mut child = Command::new("bash") - .arg("-c") - .arg(cmd_string) + let mut cmd = Command::new("bash"); + cmd.arg("-c").arg(cmd_string); + + // Drop privileges if running under sudo. This is required to avoid permission issues + // when the target command tries to access files or directories that the current user + // does not have permission to access. + if let Some((uid, gid)) = get_user_uid_gid() { + debug!("Running under sudo, dropping privileges to uid={uid}, gid={gid}"); + cmd.uid(uid).gid(gid); + } + + let mut child = cmd .spawn() .map_err(|e| anyhow!("Failed to spawn child process: {e}"))?; let root_pid = child.id() as i32; @@ -138,16 +164,21 @@ fn track_command( let writer_thread = thread::spawn(move || -> anyhow::Result<()> { let mut writer = MemtrackWriter::new(out_file)?; + let mut i = 0; while let Ok(first) = write_rx.recv() { writer.write_event(&first)?; + i += 1; // Drain any backlog in a tight loop (batching) while let Ok(ev) = write_rx.try_recv() { writer.write_event(&ev)?; + i += 1; } } writer.finish()?; + info!("Wrote {i} memtrack events to disk"); + Ok(()) }); diff --git a/crates/memtrack/src/prelude.rs b/crates/memtrack/src/prelude.rs new file mode 100644 index 00000000..7051aa84 --- /dev/null +++ b/crates/memtrack/src/prelude.rs @@ -0,0 +1,5 @@ +#[allow(unused_imports)] +pub use anyhow::{Context, Error, Result, anyhow, bail, ensure}; +pub use itertools::Itertools; +#[allow(unused_imports)] +pub use log::{debug, error, info, trace, warn};