fix: enable jemalloc-based memory control on macOS as well (#12625)

risingwavelabs · Oct 5, 2023 · e53c9f7 · e53c9f7
1 parent 8195963
commit e53c9f7
Show file tree

Hide file tree

Showing 25 changed files with 48 additions and 66 deletions.
diff --git a/.config/hakari.toml b/.config/hakari.toml
@@ -34,6 +34,9 @@ workspace-members = [
 ]
 third-party = [
     { name = "opendal" },
+    # For some reasons, tikv-jemalloc-sys would be compiled twice if being added into `workspace-hack`
+    { name = "tikv-jemalloc-sys", git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9" },
+    { name = "tikv-jemallocator", git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9" },
     # These are solely dev-dependencies. Unifying them may slow down build.
     { name = "criterion" },
     { name = "console" },

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -115,6 +115,7 @@ arrow-buffer = "46"
 arrow-flight = "46"
 arrow-select = "46"
 arrow-ord = "46"
+tikv-jemalloc-ctl = { git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9" }
 tikv-jemallocator = { git = "https://github.com/risingwavelabs/jemallocator.git", features = [
   "profiling",
   "stats",

diff --git a/src/batch/Cargo.toml b/src/batch/Cargo.toml
@@ -66,8 +66,6 @@ criterion = { workspace = true, features = ["async_tokio", "async"] }
 rand = "0.8"
 risingwave_expr_impl = { workspace = true }
 tempfile = "3"
-
-[target.'cfg(unix)'.dev-dependencies]
 tikv-jemallocator = { workspace = true }
 
 [[bench]]

diff --git a/src/batch/benches/expand.rs b/src/batch/benches/expand.rs
@@ -15,12 +15,12 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, ExpandExecutor};
-use risingwave_common::enable_jemalloc_on_unix;
+use risingwave_common::enable_jemalloc;
 use risingwave_common::types::DataType;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_expand_executor(
     column_subsets: Vec<Vec<usize>>,

diff --git a/src/batch/benches/filter.rs b/src/batch/benches/filter.rs
@@ -16,13 +16,13 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, FilterExecutor};
-use risingwave_common::enable_jemalloc_on_unix;
+use risingwave_common::enable_jemalloc;
 use risingwave_common::types::DataType;
 use risingwave_expr::expr::build_from_pretty;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_filter_executor(chunk_size: usize, chunk_num: usize) -> BoxedExecutor {
     const CHUNK_SIZE: usize = 1024;

diff --git a/src/batch/benches/hash_agg.rs b/src/batch/benches/hash_agg.rs
@@ -21,13 +21,13 @@ use risingwave_batch::task::ShutdownToken;
 use risingwave_common::catalog::{Field, Schema};
 use risingwave_common::memory::MemoryContext;
 use risingwave_common::types::DataType;
-use risingwave_common::{enable_jemalloc_on_unix, hash};
+use risingwave_common::{enable_jemalloc, hash};
 use risingwave_expr::aggregate::{AggCall, AggKind};
 use risingwave_pb::expr::{PbAggCall, PbInputRef};
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_agg_call(
     input_schema: &Schema,

diff --git a/src/batch/benches/hash_join.rs b/src/batch/benches/hash_join.rs
@@ -22,11 +22,11 @@ use risingwave_batch::task::ShutdownToken;
 use risingwave_common::catalog::schema_test_utils::field_n;
 use risingwave_common::memory::MemoryContext;
 use risingwave_common::types::DataType;
-use risingwave_common::{enable_jemalloc_on_unix, hash};
+use risingwave_common::{enable_jemalloc, hash};
 use risingwave_expr::expr::build_from_pretty;
 use utils::bench_join;
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_hash_join_executor(
     join_type: JoinType,

diff --git a/src/batch/benches/limit.rs b/src/batch/benches/limit.rs
@@ -16,12 +16,12 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, LimitExecutor};
-use risingwave_common::enable_jemalloc_on_unix;
+use risingwave_common::enable_jemalloc;
 use risingwave_common::types::DataType;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_limit_executor(
     chunk_size: usize,

diff --git a/src/batch/benches/nested_loop_join.rs b/src/batch/benches/nested_loop_join.rs
@@ -16,13 +16,13 @@ pub mod utils;
 use criterion::{criterion_group, criterion_main, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, JoinType, NestedLoopJoinExecutor};
 use risingwave_batch::task::ShutdownToken;
-use risingwave_common::enable_jemalloc_on_unix;
+use risingwave_common::enable_jemalloc;
 use risingwave_common::memory::MemoryContext;
 use risingwave_common::types::DataType;
 use risingwave_expr::expr::build_from_pretty;
 use utils::{bench_join, create_input};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_nested_loop_join_executor(
     join_type: JoinType,

diff --git a/src/batch/benches/sort.rs b/src/batch/benches/sort.rs
@@ -16,14 +16,14 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, SortExecutor};
-use risingwave_common::enable_jemalloc_on_unix;
+use risingwave_common::enable_jemalloc;
 use risingwave_common::memory::MemoryContext;
 use risingwave_common::types::DataType;
 use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_order_by_executor(
     chunk_size: usize,

diff --git a/src/batch/benches/top_n.rs b/src/batch/benches/top_n.rs
@@ -16,14 +16,14 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, TopNExecutor};
-use risingwave_common::enable_jemalloc_on_unix;
+use risingwave_common::enable_jemalloc;
 use risingwave_common::memory::MemoryContext;
 use risingwave_common::types::DataType;
 use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-enable_jemalloc_on_unix!();
+enable_jemalloc!();
 
 fn create_top_n_executor(
     chunk_size: usize,

diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml
@@ -29,6 +29,9 @@ risingwave_expr_impl = { workspace = true }
 risingwave_frontend = { workspace = true }
 risingwave_meta = { workspace = true }
 risingwave_rt = { workspace = true }
+tikv-jemallocator = { workspace = true, features = [
+    "unprefixed_malloc_on_supported_platforms",
+] }
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",
     "rt-multi-thread",
@@ -46,11 +49,6 @@ workspace-hack = { path = "../workspace-hack" }
 [target.'cfg(enable_task_local_alloc)'.dependencies]
 task_stats_alloc = { path = "../utils/task_stats_alloc" }
 
-[target.'cfg(unix)'.dependencies]
-tikv-jemallocator = { workspace = true, features = [
-    "unprefixed_malloc_on_supported_platforms",
-] }
-
 [[bin]]
 name = "frontend"
 path = "src/bin/frontend_node.rs"

diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs
@@ -24,10 +24,10 @@ use risingwave_rt::{init_risingwave_logger, main_okk, LoggerSettings};
 macro_rules! main {
     ($component:ident) => {
         #[cfg(enable_task_local_alloc)]
-        risingwave_common::enable_task_local_jemalloc_on_unix!();
+        risingwave_common::enable_task_local_jemalloc!();
 
         #[cfg(not(enable_task_local_alloc))]
-        risingwave_common::enable_jemalloc_on_unix!();
+        risingwave_common::enable_jemalloc!();
 
         #[cfg_attr(coverage, no_coverage)]
         fn main() {

diff --git a/src/cmd_all/Cargo.toml b/src/cmd_all/Cargo.toml
@@ -37,6 +37,9 @@ shell-words = "1.1.0"
 strum = "0.25"
 strum_macros = "0.25"
 tempfile = "3"
+tikv-jemallocator = { workspace = true, features = [
+    "unprefixed_malloc_on_supported_platforms",
+] }
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",
     "rt-multi-thread",
@@ -63,11 +66,6 @@ vergen = { version = "8", default-features = false, features = [
 [target.'cfg(enable_task_local_alloc)'.dependencies]
 task_stats_alloc = { path = "../utils/task_stats_alloc" }
 
-[target.'cfg(unix)'.dependencies]
-tikv-jemallocator = { workspace = true, features = [
-    "unprefixed_malloc_on_supported_platforms",
-] }
-
 [[bin]]
 name = "risingwave"
 path = "src/bin/risingwave.rs"

diff --git a/src/cmd_all/src/bin/risingwave.rs b/src/cmd_all/src/bin/risingwave.rs
@@ -31,10 +31,10 @@ use strum_macros::{Display, EnumIter, EnumString, IntoStaticStr};
 use tracing::Level;
 
 #[cfg(enable_task_local_alloc)]
-risingwave_common::enable_task_local_jemalloc_on_unix!();
+risingwave_common::enable_task_local_jemalloc!();
 
 #[cfg(not(enable_task_local_alloc))]
-risingwave_common::enable_jemalloc_on_unix!();
+risingwave_common::enable_jemalloc!();
 
 const BINARY_NAME: &str = "risingwave";
 const VERSION: &str = {

diff --git a/src/common/Cargo.toml b/src/common/Cargo.toml
@@ -85,7 +85,6 @@ strum = "0.25"
 strum_macros = "0.25"
 sysinfo = { version = "0.29", default-features = false }
 thiserror = "1"
-tikv-jemalloc-ctl = { git = "https://github.com/risingwavelabs/jemallocator.git", rev = "64a2d9" }
 tinyvec = { version = "1", features = ["rustc_1_55", "grab_spare_slice"] }
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",

diff --git a/src/common/src/jemalloc.rs b/src/common/src/jemalloc.rs
@@ -14,7 +14,7 @@
 
 /// If <https://github.com/tikv/jemallocator/issues/22> is resolved, we may inline this
 #[macro_export]
-macro_rules! enable_jemalloc_on_unix {
+macro_rules! enable_jemalloc {
     () => {
         #[cfg(unix)]
         #[global_allocator]
@@ -23,9 +23,8 @@ macro_rules! enable_jemalloc_on_unix {
 }
 
 #[macro_export]
-macro_rules! enable_task_local_jemalloc_on_unix {
+macro_rules! enable_task_local_jemalloc {
     () => {
-        #[cfg(unix)]
         #[global_allocator]
         static GLOBAL: task_stats_alloc::TaskLocalAlloc<tikv_jemallocator::Jemalloc> =
             task_stats_alloc::TaskLocalAlloc(tikv_jemallocator::Jemalloc);

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

diff --git a/src/compute/src/memory_management/mod.rs b/src/compute/src/memory_management/mod.rs
@@ -73,17 +73,10 @@ pub fn build_memory_control_policy(
 ) -> MemoryControlRef {
     use self::policy::JemallocMemoryControl;
 
-    if cfg!(target_os = "linux") {
-        Box::new(JemallocMemoryControl::new(
-            total_memory_bytes,
-            heap_profiling_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)
-    }
+    Box::new(JemallocMemoryControl::new(
+        total_memory_bytes,
+        heap_profiling_config,
+    ))
 }
 
 /// `DummyPolicy` is used for operarting systems other than Linux. It does nothing as memory control

diff --git a/src/compute/src/rpc/service/monitor_service.rs b/src/compute/src/rpc/service/monitor_service.rs
@@ -163,6 +163,7 @@ impl MonitorService for MonitorServiceImpl {
             tracing::warn!("Manually Jemalloc dump heap file failed! {:?}", e);
             Err(Status::internal(e.to_string()))
         } else {
+            tracing::info!("Manually Jemalloc dump heap file created: {}", &file_path);
             Ok(Response::new(HeapProfilingResponse {}))
         };
         response

diff --git a/src/risedevtool/src/task/compactor_service.rs b/src/risedevtool/src/task/compactor_service.rs
@@ -91,10 +91,9 @@ impl Task for CompactorService {
 
         if crate::util::is_env_set("RISEDEV_ENABLE_HEAP_PROFILE") {
             // See https://linux.die.net/man/3/jemalloc for the descriptions of profiling options
-            cmd.env(
-                "MALLOC_CONF",
-                "prof:true,lg_prof_interval:34,lg_prof_sample:19,prof_prefix:compactor",
-            );
+            let conf = "prof:true,lg_prof_interval:34,lg_prof_sample:19,prof_prefix:compactor";
+            cmd.env("_RJEM_MALLOC_CONF", conf); // prefixed for macos
+            cmd.env("MALLOC_CONF", conf); // unprefixed for linux
         }
 
         cmd.arg("--config-path")

diff --git a/src/risedevtool/src/task/compute_node_service.rs b/src/risedevtool/src/task/compute_node_service.rs
@@ -99,10 +99,9 @@ impl Task for ComputeNodeService {
 
         if crate::util::is_env_set("RISEDEV_ENABLE_HEAP_PROFILE") {
             // See https://linux.die.net/man/3/jemalloc for the descriptions of profiling options
-            cmd.env(
-                "MALLOC_CONF",
-                "prof:true,lg_prof_interval:34,lg_prof_sample:19,prof_prefix:compute-node",
-            );
+            let conf = "prof:true,lg_prof_interval:34,lg_prof_sample:19,prof_prefix:compute-node";
+            cmd.env("_RJEM_MALLOC_CONF", conf); // prefixed for macos
+            cmd.env("MALLOC_CONF", conf); // unprefixed for linux
         }
 
         cmd.arg("--config-path")

diff --git a/src/risedevtool/src/task/meta_node_service.rs b/src/risedevtool/src/task/meta_node_service.rs
@@ -184,10 +184,9 @@ impl Task for MetaNodeService {
 
         if crate::util::is_env_set("RISEDEV_ENABLE_HEAP_PROFILE") {
             // See https://linux.die.net/man/3/jemalloc for the descriptions of profiling options
-            cmd.env(
-                "MALLOC_CONF",
-                "prof:true,lg_prof_interval:32,lg_prof_sample:19,prof_prefix:meta-node",
-            );
+            let conf = "prof:true,lg_prof_interval:32,lg_prof_sample:19,prof_prefix:meta-node";
+            cmd.env("_RJEM_MALLOC_CONF", conf); // prefixed for macos
+            cmd.env("MALLOC_CONF", conf); // unprefixed for linux
         }
 
         Self::apply_command_args(&mut cmd, &self.config, HummockInMemoryStrategy::Isolated)?;

diff --git a/src/tests/simulation/Cargo.toml b/src/tests/simulation/Cargo.toml
@@ -47,14 +47,12 @@ serde_derive = "1.0.188"
 serde_json = "1.0.107"
 sqllogictest = "0.17.0"
 tempfile = "3"
+tikv-jemallocator = { workspace = true }
 tokio = { version = "0.2.23", package = "madsim-tokio" }
 tokio-postgres = "0.7"
 tokio-stream = "0.1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
-[target.'cfg(unix)'.dependencies]
-tikv-jemallocator = { workspace = true }
-
 [lints]
 workspace = true