Add jemalloc-profiling feature to enable heap profiling features.

Author: jimmygchen

Cargo.lock

@@ -23,6 +23,7 @@ lighthouse_network = { workspace = true }
 lighthouse_version = { workspace = true }
 logging = { workspace = true }
 lru = { workspace = true }
+malloc_utils = { workspace = true }
 metrics = { workspace = true }
 network = { workspace = true }
 operation_pool = { workspace = true }

beacon_node/http_api/Cargo.toml

@@ -4631,6 +4631,53 @@ pub fn serve<T: BeaconChainTypes>(
             },
         );
 
+    // POST lighthouse/malloc/prof_dump
+    let post_lighthouse_malloc_prof_dump = warp::path("lighthouse")
+        .and(warp::path("malloc"))
+        .and(warp::path("prof_dump"))
+        .and(warp::body::json())
+        .and(warp::path::end())
+        // Skip the `BeaconProcessor` for memory dumps so we can execute them as
+        // quickly as possible. Memory dumps should be uncommon and very
+        // deliberate.
+        .then(|filename: String| {
+            let dump = || {
+                let path = PathBuf::from_str(&filename).map_err(|e| {
+                    warp_utils::reject::custom_bad_request(format!(
+                        "Unable to parse {filename} as path: {e:?}"
+                    ))
+                })?;
+                if path.exists() {
+                    Err(warp_utils::reject::custom_bad_request(format!(
+                        "{filename} already exists"
+                    )))
+                } else {
+                    malloc_utils::prof_dump(&filename)
+                        .map(|()| warp::reply::json(&filename).into_response())
+                        .map_err(warp_utils::reject::custom_bad_request)
+                }
+            };
+
+            convert_rejection(dump())
+        });
+
+    // POST lighthouse/malloc/prof_active
+    let post_lighthouse_malloc_prof_active = warp::path("lighthouse")
+        .and(warp::path("malloc"))
+        .and(warp::path("prof_active"))
+        .and(warp::body::json())
+        .and(warp::path::end())
+        // Skip the `BeaconProcessor` for profiling so we can execute it as
+        // quickly as possible. Memory dumps should be uncommon and very
+        // deliberate.
+        .then(|enable: bool| {
+            let result = malloc_utils::prof_active(enable)
+                .map(|()| warp::reply::json(&enable).into_response())
+                .map_err(warp_utils::reject::custom_bad_request);
+
+            convert_rejection(result)
+        });
+
     let get_events = eth_v1
         .and(warp::path("events"))
         .and(warp::path::end())
@@ -4908,6 +4955,8 @@ pub fn serve<T: BeaconChainTypes>(
                     .uor(post_lighthouse_compaction)
                     .uor(post_lighthouse_add_peer)
                     .uor(post_lighthouse_remove_peer)
+                    .uor(post_lighthouse_malloc_prof_dump)
+                    .uor(post_lighthouse_malloc_prof_active)
                     .recover(warp_utils::reject::handle_rejection),
             ),
         )

beacon_node/http_api/src/lib.rs

@@ -7,13 +7,14 @@ edition = { workspace = true }
 [features]
 mallinfo2 = []
 jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
-jemalloc-profiling = ["tikv-jemallocator/profiling"]
+jemalloc-profiling = ["jemalloc", "tikv-jemallocator/profiling", "tikv-jemalloc-sys"]
 
 [dependencies]
 libc = "0.2.79"
 metrics = { workspace = true }
 parking_lot = { workspace = true }
 tikv-jemalloc-ctl = { version = "0.6.0", optional = true, features = ["stats"] }
+tikv-jemalloc-sys = { version = "0.6.0", optional = true }
 
 [target.'cfg(not(target_os = "linux"))'.dependencies]
 tikv-jemallocator = { version = "0.6.0", optional = true, features = ["stats"] }

common/malloc_utils/Cargo.toml

@@ -7,10 +7,13 @@
 //!
 //! A) `JEMALLOC_SYS_WITH_MALLOC_CONF` at compile-time.
 //! B) `_RJEM_MALLOC_CONF` at runtime.
+
 use metrics::{
     set_gauge, set_gauge_vec, try_create_int_gauge, try_create_int_gauge_vec, IntGauge, IntGaugeVec,
 };
+use std::ffi::{c_char, c_int};
 use std::sync::LazyLock;
+use std::{mem, ptr};
 use tikv_jemalloc_ctl::{arenas, epoch, raw, stats, Access, AsName, Error};
 
 #[global_allocator]
@@ -124,6 +127,62 @@ pub fn page_size() -> Result<usize, Error> {
     "arenas.page\0".name().read()
 }
 
+/// A convenience wrapper around `mallctl` for writing `value` to `name`.
+#[cfg(feature = "jemalloc-profiling")]
+pub unsafe fn mallctl_write<T>(name: &[u8], mut value: T) -> Result<(), c_int> {
+    // Use `tikv_jemalloc_sys::mallctl` directly since the `jemalloc_ctl::raw`
+    // functions artifically limit the `name` values.
+    let status = tikv_jemalloc_sys::mallctl(
+        name as *const _ as *const c_char,
+        ptr::null_mut(),
+        ptr::null_mut(),
+        &mut value as *mut _ as *mut _,
+        mem::size_of::<T>(),
+    );
+
+    if status == 0 {
+        Ok(())
+    } else {
+        Err(status)
+    }
+}
+
+/// Add a C-style `0x00` terminator to the string and return it as a `Vec` of
+/// bytes.
+#[allow(dead_code)]
+fn terminate_string_for_c(s: &str) -> Vec<u8> {
+    let mut terminated = vec![0x00_u8; s.len() + 1];
+    terminated[..s.len()].copy_from_slice(s.as_ref());
+    terminated
+}
+
+/// Uses `mallctl` to call `"prof.dump"`.
+///
+/// This generates a heap profile at `filename`.
+#[allow(dead_code)]
+#[cfg(feature = "jemalloc-profiling")]
+pub fn prof_dump(filename: &str) -> Result<(), String> {
+    let terminated_filename = terminate_string_for_c(filename);
+
+    unsafe {
+        mallctl_write(
+            "prof.dump\0".as_ref(),
+            terminated_filename.as_ptr() as *const c_char,
+        )
+    }
+    .map_err(|e| format!("Failed to call prof.dump on mallctl: {e:?}"))
+}
+
+/// Uses `mallctl` to call `"prof.enable"`.
+///
+/// Controls wether profile sampling is active.
+#[allow(dead_code)]
+#[cfg(feature = "jemalloc-profiling")]
+pub fn prof_active(enable: bool) -> Result<(), String> {
+    unsafe { mallctl_write("prof.active\0".as_ref(), enable) }
+        .map_err(|e| format!("Failed to call prof.active on mallctl with code {e:?}"))
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

common/malloc_utils/src/jemalloc.rs

@@ -44,6 +44,9 @@ pub use interface::*;
 mod interface {
     pub use crate::glibc::configure_glibc_malloc as configure_memory_allocator;
     pub use crate::glibc::scrape_mallinfo_metrics as scrape_allocator_metrics;
+
+    #[allow(dead_code)]
+    pub use super::prof_dump_unsupported as prof_dump;
 }
 
 #[cfg(feature = "jemalloc")]
@@ -54,6 +57,16 @@ mod interface {
     }
 
     pub use crate::jemalloc::scrape_jemalloc_metrics as scrape_allocator_metrics;
+
+    #[cfg(not(feature = "jemalloc-profiling"))]
+    pub use super::prof_dump_unsupported as prof_dump;
+    #[cfg(feature = "jemalloc-profiling")]
+    pub use crate::jemalloc::prof_dump;
+
+    #[cfg(not(feature = "jemalloc-profiling"))]
+    pub use super::prof_active_unsupported as prof_active;
+    #[cfg(feature = "jemalloc-profiling")]
+    pub use crate::jemalloc::prof_active;
 }
 
 #[cfg(all(
@@ -68,4 +81,28 @@ mod interface {
 
     #[allow(dead_code)]
     pub fn scrape_allocator_metrics() {}
+
+    #[allow(dead_code)]
+    pub use super::prof_dump_unsupported as prof_dump;
+
+    #[allow(dead_code)]
+    pub use super::prof_active_unsupported as prof_active;
+}
+
+#[allow(dead_code)]
+pub fn prof_dump_unsupported(_: &str) -> Result<(), String> {
+    Err(
+        "Profile dumps are only supported when Lighthouse is built for Linux \
+        using the `jemalloc` and `jemalloc-profiling` features."
+            .to_string(),
+    )
+}
+
+#[allow(dead_code)]
+pub fn prof_active_unsupported(_: bool) -> Result<(), String> {
+    Err(
+        "Enabling profiling is only supported when Lighthouse is built for Linux \
+        using the `jemalloc` and `jemalloc-profiling` features."
+            .to_string(),
+    )
 }

common/malloc_utils/src/lib.rs

@@ -38,6 +38,8 @@ console-subscriber = ["console-subscriber/default"]
 
 # Deprecated. This is now enabled by default on non windows targets.
 jemalloc = []
+# Enable profiling in jemalloc.
+jemalloc-profiling = ["malloc_utils/jemalloc-profiling"]
 
 [dependencies]
 account_manager = { "path" = "../account_manager" }