Make jemalloc memory profiling opt-in via --enable-memory-profiling flag

Author: ndr-ds

docker/Dockerfile

@@ -16,7 +16,7 @@ ARG binaries=
 ARG copy=${binaries:+_copy}
 ARG build_flag=--release
 ARG build_folder=release
-ARG build_features=scylladb,metrics,memory-profiling
+ARG build_features=scylladb,metrics,jemalloc
 ARG rustflags="-C force-frame-pointers=yes"
 
 FROM rust:1.74-slim-bookworm AS builder

linera-faucet/server/src/lib.rs

@@ -807,6 +807,7 @@ where
     pending_requests: Arc<Mutex<VecDeque<PendingRequest>>>,
     request_notifier: Arc<Notify>,
     max_batch_size: usize,
+    enable_memory_profiling: bool,
 }
 
 impl<C> Clone for FaucetService<C>
@@ -833,6 +834,7 @@ where
             pending_requests: Arc::clone(&self.pending_requests),
             request_notifier: Arc::clone(&self.request_notifier),
             max_batch_size: self.max_batch_size,
+            enable_memory_profiling: self.enable_memory_profiling,
         }
     }
 }
@@ -848,6 +850,7 @@ pub struct FaucetConfig {
     pub chain_listener_config: ChainListenerConfig,
     pub storage_path: PathBuf,
     pub max_batch_size: usize,
+    pub enable_memory_profiling: bool,
 }
 
 impl<C> FaucetService<C>
@@ -901,6 +904,7 @@ where
             pending_requests,
             request_notifier,
             max_batch_size: config.max_batch_size,
+            enable_memory_profiling: config.enable_memory_profiling,
         })
     }
 
@@ -937,7 +941,12 @@ where
         let index_handler = axum::routing::get(graphiql).post(Self::index_handler);
 
         #[cfg(feature = "metrics")]
-        monitoring_server::start_metrics(self.metrics_address(), cancellation_token.clone());
+        monitoring_server::start_metrics_with_profiling(
+            self.metrics_address(),
+            cancellation_token.clone(),
+            self.enable_memory_profiling,
+        )
+        .await;
 
         let app = Router::new()
             .route("/", index_handler)

linera-metrics/Cargo.toml

@@ -15,7 +15,7 @@ version.workspace = true
 workspace = true
 
 [features]
-memory-profiling = ["jemalloc_pprof"]
+jemalloc = ["jemalloc_pprof"]
 
 [dependencies]
 anyhow.workspace = true

linera-metrics/src/lib.rs

@@ -5,5 +5,5 @@
 
 pub mod monitoring_server;
 
-#[cfg(feature = "memory-profiling")]
+#[cfg(feature = "jemalloc")]
 pub mod memory_profiler;

linera-metrics/src/memory_profiler.rs

@@ -22,12 +22,33 @@ pub enum MemoryProfilerError {
     ProfCtlNotAvailable,
     #[error("another profiler is already running")]
     AnotherProfilerAlreadyRunning,
+    #[error("failed to activate jemalloc profiling: {0}")]
+    ActivationFailed(String),
 }
 
 /// Memory profiler using safe jemalloc_pprof wrapper (pull model only)
 pub struct MemoryProfiler;
 
 impl MemoryProfiler {
+    /// Activates jemalloc profiling at runtime.
+    ///
+    /// This enables sampling (prof_active) which is off by default to avoid
+    /// the libgcc DWARF unwinder livelock (jemalloc/jemalloc#2282).
+    pub async fn activate() -> Result<(), MemoryProfilerError> {
+        if let Some(prof_ctl) = PROF_CTL.as_ref() {
+            let mut prof_ctl = prof_ctl.lock().await;
+
+            prof_ctl
+                .activate()
+                .map_err(|e| MemoryProfilerError::ActivationFailed(e.to_string()))?;
+
+            info!("jemalloc memory profiling activated");
+            Ok(())
+        } else {
+            Err(MemoryProfilerError::ProfCtlNotAvailable)
+        }
+    }
+
     pub fn check_prof_ctl() -> Result<(), MemoryProfilerError> {
         // Check if jemalloc profiling is available
         if let Some(prof_ctl) = PROF_CTL.as_ref() {

linera-metrics/src/monitoring_server.rs

@@ -8,36 +8,81 @@ use tokio::net::ToSocketAddrs;
 use tokio_util::sync::CancellationToken;
 use tracing::info;
 
-#[cfg(feature = "memory-profiling")]
+#[cfg(feature = "jemalloc")]
 use crate::memory_profiler::MemoryProfiler;
 
-pub fn start_metrics(
-    address: impl ToSocketAddrs + Debug + Send + 'static,
-    shutdown_signal: CancellationToken,
-) {
-    #[cfg(feature = "memory-profiling")]
-    let app = {
-        // Try to add memory profiling endpoint
-        match MemoryProfiler::check_prof_ctl() {
-            Ok(()) => {
-                info!("Memory profiling available, enabling /debug/pprof and /debug/flamegraph endpoints");
-                Router::new()
-                    .route("/metrics", get(serve_metrics))
-                    .route("/debug/pprof", get(MemoryProfiler::heap_profile))
-                    .route("/debug/flamegraph", get(MemoryProfiler::heap_flamegraph))
-            }
+/// Whether memory profiling endpoints should be enabled on the metrics server.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MemoryProfiling {
+    Enabled,
+    Disabled,
+}
+
+impl From<bool> for MemoryProfiling {
+    fn from(enabled: bool) -> Self {
+        if enabled {
+            MemoryProfiling::Enabled
+        } else {
+            MemoryProfiling::Disabled
+        }
+    }
+}
+
+impl MemoryProfiling {
+    /// Attempts to activate jemalloc memory profiling if requested.
+    ///
+    /// Returns `MemoryProfiling::Enabled` only if activation succeeds. If the flag is
+    /// false, or jemalloc is not compiled in, or activation fails (e.g. `MALLOC_CONF` not
+    /// set), returns `MemoryProfiling::Disabled` with an appropriate warning.
+    #[cfg(feature = "jemalloc")]
+    pub async fn try_activate(requested: bool) -> Self {
+        if !requested {
+            return MemoryProfiling::Disabled;
+        }
+        match MemoryProfiler::activate().await {
+            Ok(()) => MemoryProfiling::Enabled,
             Err(e) => {
                 tracing::warn!(
-                    "Memory profiling not available: {}, serving metrics-only",
+                    "--enable-memory-profiling was passed but profiling could not be activated: {}. \
+                     Set MALLOC_CONF=\"prof:true,prof_active:false,lg_prof_sample:19\" before starting the binary.",
                     e
                 );
-                Router::new().route("/metrics", get(serve_metrics))
+                MemoryProfiling::Disabled
             }
         }
-    };
+    }
 
-    #[cfg(not(feature = "memory-profiling"))]
-    let app = Router::new().route("/metrics", get(serve_metrics));
+    /// Non-jemalloc fallback: always returns `Disabled`, warns if profiling was requested.
+    #[cfg(not(feature = "jemalloc"))]
+    pub async fn try_activate(requested: bool) -> Self {
+        if requested {
+            tracing::warn!(
+                "--enable-memory-profiling was passed but the jemalloc feature is not compiled in"
+            );
+        }
+        MemoryProfiling::Disabled
+    }
+}
+
+/// Activates memory profiling if requested and starts the metrics server.
+///
+/// This is the single entry point that handles both activation and server startup,
+/// avoiding duplication across binaries.
+pub async fn start_metrics_with_profiling(
+    address: impl ToSocketAddrs + Debug + Send + 'static,
+    shutdown_signal: CancellationToken,
+    enable_memory_profiling: bool,
+) {
+    let memory_profiling = MemoryProfiling::try_activate(enable_memory_profiling).await;
+    start_metrics(address, shutdown_signal, memory_profiling);
+}
+
+pub fn start_metrics(
+    address: impl ToSocketAddrs + Debug + Send + 'static,
+    shutdown_signal: CancellationToken,
+    memory_profiling: MemoryProfiling,
+) {
+    let app = metrics_router(memory_profiling);
 
     tokio::spawn(async move {
         let listener = tokio::net::TcpListener::bind(address)
@@ -55,6 +100,36 @@ pub fn start_metrics(
     });
 }
 
+fn metrics_router(memory_profiling: MemoryProfiling) -> Router {
+    #[cfg(feature = "jemalloc")]
+    if memory_profiling == MemoryProfiling::Enabled {
+        match MemoryProfiler::check_prof_ctl() {
+            Ok(()) => {
+                info!("Memory profiling enabled, registering /debug/pprof and /debug/flamegraph endpoints");
+                return Router::new()
+                    .route("/metrics", get(serve_metrics))
+                    .route("/debug/pprof", get(MemoryProfiler::heap_profile))
+                    .route("/debug/flamegraph", get(MemoryProfiler::heap_flamegraph));
+            }
+            Err(e) => {
+                tracing::warn!(
+                    "Memory profiling requested but not available: {}, serving metrics-only",
+                    e
+                );
+            }
+        }
+    }
+
+    #[cfg(not(feature = "jemalloc"))]
+    if memory_profiling == MemoryProfiling::Enabled {
+        tracing::warn!(
+            "Memory profiling requested but jemalloc feature is not compiled in, serving metrics-only"
+        );
+    }
+
+    Router::new().route("/metrics", get(serve_metrics))
+}
+
 async fn serve_metrics() -> Result<String, AxumError> {
     let metric_families = prometheus::gather();
     Ok(prometheus::TextEncoder::new()

linera-service/Cargo.toml

@@ -52,12 +52,10 @@ metrics = [
     "linera-faucet-server/metrics",
     "linera-metrics",
 ]
-jemalloc = ["tikv-jemallocator"]
-memory-profiling = [
-    "metrics",
-    "jemalloc",
+jemalloc = [
+    "tikv-jemallocator",
     "tikv-jemallocator/profiling",
-    "linera-metrics/memory-profiling",
+    "linera-metrics/jemalloc",
 ]
 storage-service = ["linera-storage-service"]
 opentelemetry = ["linera-rpc/opentelemetry"]

linera-service/src/cli/main.rs

@@ -8,22 +8,6 @@
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-// jemalloc configuration for memory profiling with jemalloc_pprof
-// prof:true,prof_active:true - Enable profiling from start
-// lg_prof_sample:19 - Sample every 512KB for good detail/overhead balance
-
-// Linux/other platforms: use unprefixed malloc (with unprefixed_malloc_on_supported_platforms)
-#[cfg(all(feature = "memory-profiling", not(target_os = "macos")))]
-#[allow(non_upper_case_globals)]
-#[export_name = "malloc_conf"]
-pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0";
-
-// macOS: use prefixed malloc (without unprefixed_malloc_on_supported_platforms)
-#[cfg(all(feature = "memory-profiling", target_os = "macos"))]
-#[allow(non_upper_case_globals)]
-#[export_name = "_rjem_malloc_conf"]
-pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0";
-
 mod options;
 use std::{
     collections::{BTreeMap, BTreeSet},
@@ -827,6 +811,7 @@ impl Runnable for Job {
                             monitoring_server::start_metrics(
                                 metrics_address,
                                 shutdown_notifier.clone(),
+                                monitoring_server::MemoryProfiling::Disabled,
                             );
                         }
 
@@ -1135,6 +1120,7 @@ impl Runnable for Job {
                             monitoring_server::start_metrics(
                                 metrics_address,
                                 shutdown_notifier.clone(),
+                                monitoring_server::MemoryProfiling::Disabled,
                             );
                         }
 
@@ -1319,6 +1305,7 @@ impl Runnable for Job {
                     query_cache_size,
                     query_subscriptions,
                     cancellation_token.clone(),
+                    options.enable_memory_profiling(),
                 );
                 service.run(cancellation_token, command_receiver).await?;
             }
@@ -1351,6 +1338,7 @@ impl Runnable for Job {
                         u64::try_from(et.timestamp_micros()).expect("End timestamp before 1970");
                     Timestamp::from(micros)
                 });
+
                 let config = FaucetConfig {
                     port,
                     #[cfg(with_metrics)]
@@ -1362,6 +1350,7 @@ impl Runnable for Job {
                     chain_listener_config: config,
                     storage_path,
                     max_batch_size,
+                    enable_memory_profiling: options.enable_memory_profiling(),
                 };
                 let faucet = FaucetService::new(config, context).await?;
                 let cancellation_token = CancellationToken::new();

linera-service/src/cli/options.rs

@@ -90,6 +90,11 @@ pub struct Options {
     )]
     pub execution_state_cache_size: usize,
 
+    /// Enable jemalloc memory profiling endpoints on the metrics server.
+    #[cfg(feature = "jemalloc")]
+    #[arg(long, env = "LINERA_ENABLE_MEMORY_PROFILING")]
+    pub enable_memory_profiling: bool,
+
     /// Subcommand.
     #[command(subcommand)]
     pub command: ClientCommand,
@@ -100,6 +105,17 @@ impl Options {
         <Options as clap::Parser>::parse()
     }
 
+    pub fn enable_memory_profiling(&self) -> bool {
+        #[cfg(feature = "jemalloc")]
+        {
+            self.enable_memory_profiling
+        }
+        #[cfg(not(feature = "jemalloc"))]
+        {
+            false
+        }
+    }
+
     pub async fn create_client_context<S, Si>(
         &self,
         storage: S,