feat(build): use tikv-jemalloc only on Linux

docs/memory-profiling.md

@@ -1,12 +1,14 @@
 # Memory (Heap) Profiling Guide
 
+Note that the content below is Linux-exclusive.
+
 ## What is Heap Profile?
 
 A heap profiler records the **stack trace of the allocation** of each **live** object, so it’s possible that function A allocates something and then hand over it to struct B, in this case, the allocation will still be counted on A.
 
 ## Internals
 
-RisingWave uses [tikv-jemallocator](https://crates.io/crates/tikv-jemallocator), which is a Rust wrapper of [jemalloc](https://github.com/jemalloc/jemalloc), as its memory allocator.
+RisingWave uses [tikv-jemallocator](https://crates.io/crates/tikv-jemallocator) on Linux, which is a Rust wrapper of [jemalloc](https://github.com/jemalloc/jemalloc), as its memory allocator. On other platforms, RisingWave uses the default allocator.
 
 Luckily, jemalloc provides built-in profiling support ([official wiki](https://github.com/jemalloc/jemalloc/wiki/Use-Case%3A-Heap-Profiling)). jemallocator exposes the feature via a cargo feature ‘profiling’. [Here](https://gist.github.com/ordian/928dc2bd45022cddd547528f64db9174) is a simple guide to profiling with jemallocator.
 

src/batch/Cargo.toml

@@ -71,6 +71,8 @@ workspace-hack = { path = "../workspace-hack" }
 criterion = { version = "0.4", features = ["async_tokio", "async"] }
 rand = "0.8"
 tempfile = "3"
+
+[target.'cfg(target_os = "linux")'.dev-dependencies]
 tikv-jemallocator = "0.5"
 
 [[bench]]

src/batch/benches/expand.rs

@@ -15,13 +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_linux;
 use risingwave_common::types::DataType;
-use tikv_jemallocator::Jemalloc;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_expand_executor(
     column_subsets: Vec<Vec<usize>>,

src/batch/benches/filter.rs

@@ -16,6 +16,7 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, FilterExecutor};
+use risingwave_common::enable_jemalloc_on_linux;
 use risingwave_common::types::{DataType, ScalarImpl};
 use risingwave_common::util::value_encoding::serialize_datum;
 use risingwave_expr::expr::build_from_prost;
@@ -26,12 +27,10 @@ use risingwave_pb::expr::expr_node::Type::{
     ConstantValue as TConstValue, Equal, InputRef, Modulus,
 };
 use risingwave_pb::expr::{ExprNode, FunctionCall, InputRefExpr};
-use tikv_jemallocator::Jemalloc;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_filter_executor(chunk_size: usize, chunk_num: usize) -> BoxedExecutor {
     const CHUNK_SIZE: usize = 1024;

src/batch/benches/hash_agg.rs

@@ -17,18 +17,16 @@ use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criteri
 use itertools::Itertools;
 use risingwave_batch::executor::{BoxedExecutor, HashAggExecutor};
 use risingwave_common::catalog::{Field, Schema};
-use risingwave_common::hash;
 use risingwave_common::types::DataType;
+use risingwave_common::{enable_jemalloc_on_linux, hash};
 use risingwave_expr::expr::AggKind;
 use risingwave_expr::vector_op::agg::AggStateFactory;
 use risingwave_pb::expr::agg_call::Arg;
 use risingwave_pb::expr::{AggCall, InputRefExpr};
-use tikv_jemallocator::Jemalloc;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_agg_call(
     input_schema: &Schema,

src/batch/benches/hash_join.rs

@@ -19,9 +19,9 @@ use risingwave_batch::executor::hash_join::HashJoinExecutor;
 use risingwave_batch::executor::test_utils::{gen_projected_data, MockExecutor};
 use risingwave_batch::executor::{BoxedExecutor, JoinType};
 use risingwave_common::catalog::schema_test_utils::field_n;
-use risingwave_common::hash;
 use risingwave_common::types::{DataType, ScalarImpl};
 use risingwave_common::util::value_encoding::serialize_datum;
+use risingwave_common::{enable_jemalloc_on_linux, hash};
 use risingwave_expr::expr::build_from_prost;
 use risingwave_pb::data::data_type::TypeName;
 use risingwave_pb::data::Datum as ProstDatum;
@@ -30,11 +30,9 @@ use risingwave_pb::expr::expr_node::Type::{
     ConstantValue as TConstValue, GreaterThan, InputRef, Modulus,
 };
 use risingwave_pb::expr::{ExprNode, FunctionCall, InputRefExpr};
-use tikv_jemallocator::Jemalloc;
 use utils::bench_join;
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_hash_join_executor(
     join_type: JoinType,

src/batch/benches/limit.rs

@@ -16,13 +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_linux;
 use risingwave_common::types::DataType;
-use tikv_jemallocator::Jemalloc;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_limit_executor(
     chunk_size: usize,

src/batch/benches/nested_loop_join.rs

@@ -15,6 +15,7 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, JoinType, NestedLoopJoinExecutor};
+use risingwave_common::enable_jemalloc_on_linux;
 use risingwave_common::types::{DataType, ScalarImpl};
 use risingwave_common::util::value_encoding::serialize_datum;
 use risingwave_expr::expr::build_from_prost;
@@ -25,11 +26,9 @@ use risingwave_pb::expr::expr_node::Type::{
     ConstantValue as TConstValue, Equal, InputRef, Modulus,
 };
 use risingwave_pb::expr::{ExprNode, FunctionCall, InputRefExpr};
-use tikv_jemallocator::Jemalloc;
 use utils::{bench_join, create_input};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_nested_loop_join_executor(
     join_type: JoinType,

src/batch/benches/sort.rs

@@ -16,14 +16,13 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, SortExecutor};
+use risingwave_common::enable_jemalloc_on_linux;
 use risingwave_common::types::DataType;
 use risingwave_common::util::sort_util::{OrderPair, OrderType};
-use tikv_jemallocator::Jemalloc;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_order_by_executor(
     chunk_size: usize,

src/batch/benches/sort_merge_join.rs

@@ -18,13 +18,12 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use risingwave_batch::executor::test_utils::{gen_sorted_data, MockExecutor};
 use risingwave_batch::executor::{BoxedExecutor, JoinType, SortMergeJoinExecutor};
 use risingwave_common::catalog::schema_test_utils::field_n;
+use risingwave_common::enable_jemalloc_on_linux;
 use risingwave_common::types::DataType;
 use risingwave_common::util::sort_util::OrderType;
-use tikv_jemallocator::Jemalloc;
 use utils::bench_join;
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_sort_merge_join_executor(
     join_type: JoinType,

src/batch/benches/top_n.rs

@@ -16,14 +16,13 @@ pub mod utils;
 
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
 use risingwave_batch::executor::{BoxedExecutor, TopNExecutor};
+use risingwave_common::enable_jemalloc_on_linux;
 use risingwave_common::types::DataType;
 use risingwave_common::util::sort_util::{OrderPair, OrderType};
-use tikv_jemallocator::Jemalloc;
 use tokio::runtime::Runtime;
 use utils::{create_input, execute_executor};
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 fn create_top_n_executor(
     chunk_size: usize,

src/cmd/Cargo.toml

@@ -15,6 +15,7 @@ static-log-level = ["workspace-config/enable-static-log-level"]
 anyhow = "1"
 clap = { version = "3", features = ["derive"] }
 log = { version = "0.4" }
+risingwave_common = { path = "../common" }
 risingwave_compaction_test = { path = "../tests/compaction_test" }
 risingwave_compactor = { path = "../storage/compactor" }
 risingwave_compute = { path = "../compute" }
@@ -23,7 +24,6 @@ risingwave_frontend = { path = "../frontend" }
 risingwave_meta = { path = "../meta" }
 risingwave_rt = { path = "../utils/runtime" }
 task_stats_alloc = { path = "../utils/task_stats_alloc" }
-tikv-jemallocator = { version = "0.5", features = ["profiling", "stats"] }
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",
     "rt-multi-thread",
@@ -39,6 +39,9 @@ tracing = { version = "0.1" }
 workspace-config = { path = "../utils/workspace-config", optional = true }
 workspace-hack = { path = "../workspace-hack" }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+tikv-jemallocator = { version = "0.5", features = ["profiling", "stats"] }
+
 [[bin]]
 name = "frontend"
 path = "src/bin/frontend_node.rs"

src/cmd/src/bin/compactor.rs

@@ -14,10 +14,9 @@
 
 #![cfg_attr(coverage, feature(no_coverage))]
 
-use tikv_jemallocator::Jemalloc;
+use risingwave_common::enable_jemalloc_on_linux;
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 #[cfg_attr(coverage, no_coverage)]
 fn main() {

src/cmd/src/bin/compute_node.rs

@@ -14,11 +14,9 @@
 
 #![cfg_attr(coverage, feature(no_coverage))]
 
-use task_stats_alloc::TaskLocalAlloc;
-use tikv_jemallocator::Jemalloc;
+use risingwave_common::enable_task_local_jemalloc_on_linux;
 
-#[global_allocator]
-static GLOBAL: TaskLocalAlloc<Jemalloc> = TaskLocalAlloc(Jemalloc);
+enable_task_local_jemalloc_on_linux!();
 
 #[cfg_attr(coverage, no_coverage)]
 fn main() {

src/cmd/src/bin/ctl.rs

@@ -15,10 +15,9 @@
 #![cfg_attr(coverage, feature(no_coverage))]
 
 use anyhow::Result;
-use tikv_jemallocator::Jemalloc;
+use risingwave_common::enable_jemalloc_on_linux;
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 #[cfg_attr(coverage, no_coverage)]
 fn main() -> Result<()> {

src/cmd/src/bin/frontend_node.rs

@@ -14,10 +14,9 @@
 
 #![cfg_attr(coverage, feature(no_coverage))]
 
-use tikv_jemallocator::Jemalloc;
+use risingwave_common::enable_jemalloc_on_linux;
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 #[cfg_attr(coverage, no_coverage)]
 fn main() {

src/cmd/src/bin/meta_node.rs

@@ -14,10 +14,9 @@
 
 #![cfg_attr(coverage, feature(no_coverage))]
 
-use tikv_jemallocator::Jemalloc;
+use risingwave_common::enable_jemalloc_on_linux;
 
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+enable_jemalloc_on_linux!();
 
 #[cfg_attr(coverage, no_coverage)]
 fn main() {

src/cmd_all/src/bin/risingwave.rs

@@ -15,18 +15,15 @@
 #![cfg_attr(coverage, feature(no_coverage))]
 #![feature(let_chains)]
 
-use task_stats_alloc::TaskLocalAlloc;
-use tikv_jemallocator::Jemalloc;
-
-#[global_allocator]
-static GLOBAL: TaskLocalAlloc<Jemalloc> = TaskLocalAlloc(Jemalloc);
-
 use std::collections::HashMap;
 use std::env;
 
 use anyhow::{bail, Result};
 use clap::StructOpt;
 use risingwave_cmd_all::playground;
+use risingwave_common::enable_task_local_jemalloc_on_linux;
+
+enable_task_local_jemalloc_on_linux!();
 
 type RwFns = HashMap<&'static str, Box<dyn Fn(Vec<String>) -> Result<()>>>;
 

src/common/src/jemalloc.rs

@@ -0,0 +1,33 @@
+// Copyright 2023 Singularity Data
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// If <https://github.com/tikv/jemallocator/issues/22> is resolved, we may inline this
+#[macro_export]
+macro_rules! enable_jemalloc_on_linux {
+    () => {
+        #[cfg(target_os = "linux")]
+        #[global_allocator]
+        static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+    };
+}
+
+#[macro_export]
+macro_rules! enable_task_local_jemalloc_on_linux {
+    () => {
+        #[cfg(target_os = "linux")]
+        #[global_allocator]
+        static GLOBAL: task_stats_alloc::TaskLocalAlloc<tikv_jemallocator::Jemalloc> =
+            task_stats_alloc::TaskLocalAlloc(tikv_jemallocator::Jemalloc);
+    };
+}

src/common/src/lib.rs

@@ -34,6 +34,8 @@
 #![feature(array_chunks)]
 #![allow(incomplete_features)]
 
+#[macro_use]
+pub mod jemalloc;
 #[macro_use]
 pub mod error;
 #[macro_use]

src/compute/Cargo.toml

@@ -53,7 +53,6 @@ smallvec = "1"
 static_assertions = "1"
 sysinfo = "0.26"
 thiserror = "1"
-tikv-jemalloc-ctl = "0.5"
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",
     "rt-multi-thread",
@@ -71,6 +70,9 @@ tracing = "0.1"
 twox-hash = "1"
 url = "2"
 
+[target.'cfg(target_os = "linux")'.dependencies]
+tikv-jemalloc-ctl = "0.5"
+
 [target.'cfg(not(madsim))'.dependencies]
 workspace-hack = { path = "../workspace-hack" }