make pprof optional

kwai · Jan 24, 2025 · dfa9d2f · dfa9d2f
1 parent 9beeb84
commit dfa9d2f
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 20 deletions.
diff --git a/native-engine/blaze/Cargo.toml b/native-engine/blaze/Cargo.toml
@@ -10,7 +10,11 @@ crate-type = ["cdylib"]
 [features]
 default = ["tokio/rt-multi-thread", "jemalloc"]
 
-jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemalloc-sys", "dep:tikv-jemallocator", "dep:jemalloc_pprof", "dep:pprof"]
+jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemalloc-sys", "dep:tikv-jemallocator"]
+
+jemalloc-pprof = ["dep:jemalloc_pprof", "dep:pprof", "http-service"]
+
+http-service = []
 
 [dependencies]
 arrow = { workspace = true }
@@ -53,4 +57,4 @@ optional = true
 [dependencies.pprof]
 version = "0.14.0"
 features = ["flamegraph", "protobuf-codec", "protobuf"]
-optional = true
+optional = true
diff --git a/native-engine/blaze/src/alloc.rs b/native-engine/blaze/src/alloc.rs
@@ -14,9 +14,6 @@ use std::{
 #[cfg_attr(not(windows), global_allocator)]
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-#[cfg(not(feature = "jemalloc"))]
-static GLOBAL: std::alloc::System = std::alloc::System;
-
 // only used for debugging
 //
 // #[global_allocator]

diff --git a/native-engine/blaze/src/exec.rs b/native-engine/blaze/src/exec.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{error::Error, sync::Arc};
+use std::sync::Arc;
 
 use blaze_jni_bridge::{
     conf::{DoubleConf, IntConf},
@@ -33,14 +33,9 @@ use jni::{
     objects::{JClass, JObject},
     JNIEnv,
 };
-use once_cell::sync::{Lazy, OnceCell};
+use once_cell::sync::OnceCell;
 
-use crate::{
-    handle_unwinded_scope,
-    http::{HttpService, HTTP_SERVICE},
-    logging::init_logging,
-    rt::NativeExecutionRuntime,
-};
+use crate::{handle_unwinded_scope, logging::init_logging, rt::NativeExecutionRuntime};
 
 #[allow(non_snake_case)]
 #[no_mangle]
@@ -54,10 +49,14 @@ pub extern "system" fn Java_org_apache_spark_sql_blaze_JniBridge_callNative(
         static SESSION: OnceCell<SessionContext> = OnceCell::new();
         static INIT: OnceCell<()> = OnceCell::new();
 
-        let _ = HTTP_SERVICE.get_or_try_init(|| {
-            eprintln!("initializing http service...");
-            Ok::<HttpService, DataFusionError>(HttpService::init())
-        });
+        #[cfg(feature = "http-service")]
+        {
+            use crate::http::{HttpService, HTTP_SERVICE};
+            let _ = HTTP_SERVICE.get_or_try_init(|| {
+                eprintln!("initializing http service...");
+                Ok::<HttpService, DataFusionError>(HttpService::init())
+            });
+        }
 
         INIT.get_or_try_init(|| {
             // logging is not initialized at this moment

diff --git a/native-engine/blaze/src/http/mod.rs b/native-engine/blaze/src/http/mod.rs
@@ -1,12 +1,11 @@
+#[cfg(feature = "jemalloc-pprof")]
 mod pprof;
 
 use std::sync::Mutex;
 
 use once_cell::sync::OnceCell;
 use poem::{listener::TcpListener, Route, RouteMethod, Server};
 
-use crate::http::pprof::PProfHandler;
-
 pub static HTTP_SERVICE: OnceCell<HttpService> = OnceCell::new();
 
 pub trait Handler {
@@ -77,7 +76,11 @@ pub struct HttpService;
 impl HttpService {
     pub fn init() -> Self {
         let server = DefaultHTTPServer::new();
-        server.register_handler(PProfHandler::default());
+        #[cfg(feature = "jemalloc-pprof")]
+        {
+            use crate::http::pprof::PProfHandler;
+            server.register_handler(PProfHandler::default());
+        }
         server.start();
         Self
     }