chore: use cfg! to instead of #cfg[] for jemalloc control policy (#12307

)

src/compute/Cargo.toml

@@ -39,7 +39,7 @@ risingwave_storage = { workspace = true }
 risingwave_stream = { workspace = true }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
-
+tikv-jemalloc-ctl = { git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9" }
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",
     "rt-multi-thread",
@@ -54,8 +54,6 @@ tonic = { workspace = true }
 tower = { version = "0.4", features = ["util", "load-shed"] }
 tracing = "0.1"
 
-[target.'cfg(target_os = "linux")'.dependencies]
-tikv-jemalloc-ctl = { git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9" }
 [target.'cfg(not(madsim))'.dependencies]
 workspace-hack = { path = "../workspace-hack" }
 

src/compute/src/memory_management/memory_manager.rs

@@ -50,8 +50,7 @@ impl GlobalMemoryManager {
         auto_dump_heap_profile_config: AutoDumpHeapProfileConfig,
     ) -> Arc<Self> {
         let memory_control_policy =
-            build_memory_control_policy(total_memory_bytes, auto_dump_heap_profile_config.clone())
-                .unwrap();
+            build_memory_control_policy(total_memory_bytes, auto_dump_heap_profile_config.clone());
         tracing::info!("memory control policy: {:?}", &memory_control_policy);
 
         if auto_dump_heap_profile_config.enabled {

src/compute/src/memory_management/mod.rs

@@ -16,15 +16,13 @@ pub mod memory_manager;
 
 // Only enable the non-trivial policies on Linux as it relies on statistics from `jemalloc-ctl`
 // which might be inaccurate on other platforms.
-#[cfg(target_os = "linux")]
 pub mod policy;
 
 use std::sync::atomic::AtomicU64;
 use std::sync::Arc;
 
 use risingwave_batch::task::BatchManager;
 use risingwave_common::config::{AutoDumpHeapProfileConfig, StorageConfig, StorageMemoryConfig};
-use risingwave_common::error::Result;
 use risingwave_common::util::pretty_bytes::convert;
 use risingwave_stream::task::LocalStreamManager;
 
@@ -69,28 +67,23 @@ pub trait MemoryControl: Send + Sync + std::fmt::Debug {
     ) -> MemoryControlStats;
 }
 
-#[cfg(target_os = "linux")]
 pub fn build_memory_control_policy(
     total_memory_bytes: usize,
     auto_dump_heap_profile_config: AutoDumpHeapProfileConfig,
-) -> Result<MemoryControlRef> {
+) -> MemoryControlRef {
     use self::policy::JemallocMemoryControl;
 
-    Ok(Box::new(JemallocMemoryControl::new(
-        total_memory_bytes,
-        auto_dump_heap_profile_config,
-    )))
-}
-
-#[cfg(not(target_os = "linux"))]
-pub fn build_memory_control_policy(
-    _total_memory_bytes: usize,
-    _auto_dump_heap_profile_config: AutoDumpHeapProfileConfig,
-) -> Result<MemoryControlRef> {
-    // We disable memory control on operating systems other than Linux now because jemalloc
-    // stats do not work well.
-    tracing::warn!("memory control is only enabled on Linux now");
-    Ok(Box::new(DummyPolicy))
+    if cfg!(target_os = "linux") {
+        Box::new(JemallocMemoryControl::new(
+            total_memory_bytes,
+            auto_dump_heap_profile_config,
+        ))
+    } else {
+        // We disable memory control on operating systems other than Linux now because jemalloc
+        // stats do not work well.
+        tracing::warn!("memory control is only enabled on Linux now");
+        Box::new(DummyPolicy)
+    }
 }
 
 /// `DummyPolicy` is used for operarting systems other than Linux. It does nothing as memory control

src/compute/src/rpc/service/monitor_service.rs

@@ -105,7 +105,6 @@ impl MonitorService for MonitorServiceImpl {
         }
     }
 
-    #[cfg(target_os = "linux")]
     #[cfg_attr(coverage, no_coverage)]
     async fn heap_profiling(
         &self,
@@ -117,6 +116,12 @@ impl MonitorService for MonitorServiceImpl {
 
         use tikv_jemalloc_ctl;
 
+        if !cfg!(target_os = "linux") {
+            return Err(Status::unimplemented(
+                "heap profiling is only implemented on Linux",
+            ));
+        }
+
         if !tikv_jemalloc_ctl::opt::prof::read().unwrap() {
             return Err(Status::failed_precondition(
                 "Jemalloc profiling is not enabled on the node. Try start the node with `MALLOC_CONF=prof:true`",
@@ -147,17 +152,6 @@ impl MonitorService for MonitorServiceImpl {
         let _ = unsafe { Box::from_raw(file_path_ptr) };
         response
     }
-
-    #[cfg(not(target_os = "linux"))]
-    #[cfg_attr(coverage, no_coverage)]
-    async fn heap_profiling(
-        &self,
-        _request: Request<HeapProfilingRequest>,
-    ) -> Result<Response<HeapProfilingResponse>, Status> {
-        Err(Status::unimplemented(
-            "heap profiling is only implemented on Linux",
-        ))
-    }
 }
 
 pub use grpc_middleware::*;

src/workspace-hack/Cargo.toml

@@ -94,6 +94,7 @@ serde_json = { version = "1", features = ["alloc"] }
 serde_with = { version = "3", features = ["json"] }
 smallvec = { version = "1", default-features = false, features = ["serde", "union", "write"] }
 subtle = { version = "2" }
+tikv-jemalloc-sys = { git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9", features = ["profiling", "stats"] }
 time = { version = "0.3", features = ["local-offset", "macros", "serde-well-known"] }
 tinyvec = { version = "1", features = ["alloc", "grab_spare_slice", "rustc_1_55"] }
 tokio = { version = "1", features = ["full", "stats", "tracing"] }
@@ -195,6 +196,7 @@ smallvec = { version = "1", default-features = false, features = ["serde", "unio
 subtle = { version = "2" }
 syn-dff4ba8e3ae991db = { package = "syn", version = "1", features = ["extra-traits", "full", "visit", "visit-mut"] }
 syn-f595c2ba2a3f28df = { package = "syn", version = "2", features = ["extra-traits", "fold", "full", "visit", "visit-mut"] }
+tikv-jemalloc-sys = { git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9", features = ["profiling", "stats"] }
 time = { version = "0.3", features = ["local-offset", "macros", "serde-well-known"] }
 time-macros = { version = "0.2", default-features = false, features = ["formatting", "parsing", "serde"] }
 tinyvec = { version = "1", features = ["alloc", "grab_spare_slice", "rustc_1_55"] }