refactor(memtrack): use standard allocation API across allocators

Author: not-matthias

crates/memtrack/src/ebpf/memtrack.rs

@@ -262,40 +262,6 @@ impl MemtrackBpf {
     // 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, "posix_memalign");
-        self.try_attach_memalign(lib_path, "memalign");
-        Ok(())
-    }
-
-    /// Attach C++ operator new/delete probes.
-    /// These are mangled C++ symbols that wrap the underlying allocator.
-    /// C++ operators have identical signatures to malloc/free, so we reuse those handlers.
-    pub fn attach_libcpp_probes(&mut self, lib_path: &Path) -> Result<()> {
-        self.try_attach_malloc(lib_path, "_Znwm"); // operator new(size_t)
-        self.try_attach_malloc(lib_path, "_Znam"); // operator new[](size_t)
-        self.try_attach_malloc(lib_path, "_ZnwmSt11align_val_t"); // operator new(size_t, std::align_val_t)
-        self.try_attach_malloc(lib_path, "_ZnamSt11align_val_t"); // operator new[](size_t, std::align_val_t)
-        self.try_attach_free(lib_path, "_ZdlPv"); // operator delete(void*)
-        self.try_attach_free(lib_path, "_ZdaPv"); // operator delete[](void*)
-        self.try_attach_free(lib_path, "_ZdlPvm"); // operator delete(void*, size_t) - C++14 sized delete
-        self.try_attach_free(lib_path, "_ZdaPvm"); // operator delete[](void*, size_t) - C++14 sized delete
-        self.try_attach_free(lib_path, "_ZdlPvSt11align_val_t"); // operator delete(void*, std::align_val_t)
-        self.try_attach_free(lib_path, "_ZdaPvSt11align_val_t"); // operator delete[](void*, std::align_val_t)
-        self.try_attach_free(lib_path, "_ZdlPvmSt11align_val_t"); // operator delete(void*, size_t, std::align_val_t)
-        self.try_attach_free(lib_path, "_ZdaPvmSt11align_val_t"); // operator delete[](void*, size_t, std::align_val_t)
-
-        Ok(())
-    }
-
     /// Attach probes for a specific allocator kind.
     /// This attaches both standard probes (if the allocator exports them) and
     /// allocator-specific prefixed probes.
@@ -316,22 +282,77 @@ impl MemtrackBpf {
                 self.attach_libcpp_probes(lib_path)
             }
             AllocatorKind::Jemalloc => {
-                // Try standard names (jemalloc may export these as drop-in replacements)
-                let _ = self.attach_libc_probes(lib_path);
                 // Try C++ operators (jemalloc exports these for C++ programs)
                 let _ = self.attach_libcpp_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);
                 // Try C++ operators (mimalloc exports these for C++ programs)
                 let _ = self.attach_libcpp_probes(lib_path);
                 self.attach_mimalloc_probes(lib_path)
             }
         }
     }
 
+    fn attach_standard_probes(
+        &mut self,
+        lib_path: &Path,
+        prefixes: &[&str],
+        suffixes: &[&str],
+    ) -> Result<()> {
+        // Always include "" to capture the basic case
+        let prefixes_with_base: Vec<&str> = std::iter::once("")
+            .chain(prefixes.iter().copied())
+            .unique()
+            .collect();
+
+        let suffixes_with_base: Vec<&str> = std::iter::once("")
+            .chain(suffixes.iter().copied())
+            .unique()
+            .collect();
+
+        for prefix in &prefixes_with_base {
+            for suffix in &suffixes_with_base {
+                self.try_attach_malloc(lib_path, &format!("{prefix}malloc{suffix}"));
+                self.try_attach_calloc(lib_path, &format!("{prefix}calloc{suffix}"));
+                self.try_attach_realloc(lib_path, &format!("{prefix}realloc{suffix}"));
+                self.try_attach_aligned_alloc(lib_path, &format!("{prefix}aligned_alloc{suffix}"));
+                self.try_attach_memalign(lib_path, &format!("{prefix}memalign{suffix}"));
+                self.try_attach_memalign(lib_path, &format!("{prefix}posix_memalign{suffix}"));
+                self.try_attach_free(lib_path, &format!("{prefix}free{suffix}"));
+            }
+        }
+
+        Ok(())
+    }
+
+    /// 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.
+    fn attach_libc_probes(&mut self, lib_path: &Path) -> Result<()> {
+        self.attach_standard_probes(lib_path, &[], &[])
+    }
+
+    /// Attach C++ operator new/delete probes.
+    /// These are mangled C++ symbols that wrap the underlying allocator.
+    /// C++ operators have identical signatures to malloc/free, so we reuse those handlers.
+    fn attach_libcpp_probes(&mut self, lib_path: &Path) -> Result<()> {
+        self.try_attach_malloc(lib_path, "_Znwm"); // operator new(size_t)
+        self.try_attach_malloc(lib_path, "_Znam"); // operator new[](size_t)
+        self.try_attach_malloc(lib_path, "_ZnwmSt11align_val_t"); // operator new(size_t, std::align_val_t)
+        self.try_attach_malloc(lib_path, "_ZnamSt11align_val_t"); // operator new[](size_t, std::align_val_t)
+        self.try_attach_free(lib_path, "_ZdlPv"); // operator delete(void*)
+        self.try_attach_free(lib_path, "_ZdaPv"); // operator delete[](void*)
+        self.try_attach_free(lib_path, "_ZdlPvm"); // operator delete(void*, size_t) - C++14 sized delete
+        self.try_attach_free(lib_path, "_ZdaPvm"); // operator delete[](void*, size_t) - C++14 sized delete
+        self.try_attach_free(lib_path, "_ZdlPvSt11align_val_t"); // operator delete(void*, std::align_val_t)
+        self.try_attach_free(lib_path, "_ZdaPvSt11align_val_t"); // operator delete[](void*, std::align_val_t)
+        self.try_attach_free(lib_path, "_ZdlPvmSt11align_val_t"); // operator delete(void*, size_t, std::align_val_t)
+        self.try_attach_free(lib_path, "_ZdaPvmSt11align_val_t"); // operator delete[](void*, size_t, std::align_val_t)
+
+        Ok(())
+    }
+
     /// 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:
@@ -340,36 +361,23 @@ impl MemtrackBpf {
         // - rust_dealloc: _rjem_sdallocx
         // - rust_realloc: _rjem_realloc / _rjem_rallocx
 
-        // je_*_default API (C++ with static linking)
-        self.try_attach_malloc(lib_path, "je_malloc_default");
-        self.try_attach_malloc(lib_path, "je_mallocx_default");
-        self.try_attach_free(lib_path, "je_free_default");
-        self.try_attach_free(lib_path, "je_sdallocx_default");
-        self.try_attach_realloc(lib_path, "je_realloc_default");
-        self.try_attach_realloc(lib_path, "je_rallocx_default");
-        self.try_attach_calloc(lib_path, "je_calloc_default");
-
         // je_* API (internal jemalloc functions, static linking)
-        self.try_attach_malloc(lib_path, "je_malloc");
-        self.try_attach_malloc(lib_path, "je_mallocx");
-        self.try_attach_calloc(lib_path, "je_calloc");
-        self.try_attach_realloc(lib_path, "je_realloc");
-        self.try_attach_realloc(lib_path, "je_rallocx");
-        self.try_attach_aligned_alloc(lib_path, "je_aligned_alloc");
-        self.try_attach_memalign(lib_path, "je_memalign");
-        self.try_attach_free(lib_path, "je_free");
-        self.try_attach_free(lib_path, "je_sdallocx");
-
         // _rjem_* API (Rust jemalloc crate, dynamic linking)
-        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");
+        let prefixes = ["je_", "_rjem_"];
+        let suffixes = ["", "_default"];
+
+        self.attach_standard_probes(lib_path, &prefixes, &suffixes)?;
+
+        // Non-standard API that has an additional flag parameter
+        // See: https://jemalloc.net/jemalloc.3.html
+        for prefix in prefixes {
+            for suffix in suffixes {
+                self.try_attach_malloc(lib_path, &format!("{prefix}mallocx{suffix}"));
+                self.try_attach_realloc(lib_path, &format!("{prefix}rallocx{suffix}"));
+                self.try_attach_free(lib_path, &format!("{prefix}dallocx{suffix}"));
+                self.try_attach_free(lib_path, &format!("{prefix}sdallocx{suffix}"));
+            }
+        }
 
         Ok(())
     }
@@ -382,16 +390,10 @@ impl MemtrackBpf {
         // - 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");
+        self.attach_standard_probes(lib_path, &["mi_"], &[])?;
 
         // Zero-initialized and aligned variants
+        self.try_attach_calloc(lib_path, "mi_malloc_aligned");
         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");