chore(allocator): remove things related to task_local_alloc

.cargo/config.toml

@@ -31,13 +31,7 @@ rustflags = [
 
 # Flags for all targets.
 [target.'cfg(all())']
-rustflags = [
-    "--cfg",
-    "tokio_unstable",
-    # Uncomment the following two lines to enable `TaskLocalAlloc`.
-    # "--cfg",
-    # "enable_task_local_alloc",
-]
+rustflags = ["--cfg", "tokio_unstable"]
 
 [unstable]
 lints = true

src/batch/Cargo.toml

@@ -62,9 +62,6 @@ tokio-util = { workspace = true }
 tonic = { workspace = true }
 tracing = "0.1"
 
-[target.'cfg(enable_task_local_alloc)'.dependencies]
-task_stats_alloc = { path = "../utils/task_stats_alloc" }
-
 [target.'cfg(not(madsim))'.dependencies]
 workspace-hack = { path = "../workspace-hack" }
 

src/batch/src/task/task_execution.rs

@@ -13,12 +13,8 @@
 // limitations under the License.
 
 use std::fmt::{Debug, Formatter};
-#[cfg(enable_task_local_alloc)]
-use std::future::Future;
 use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
-#[cfg(enable_task_local_alloc)]
-use std::time::Duration;
 
 use anyhow::Context;
 use futures::StreamExt;
@@ -51,52 +47,6 @@ use crate::task::BatchTaskContext;
 // Now we will only at most have 2 status for each status channel. Running -> Failed or Finished.
 pub const TASK_STATUS_BUFFER_SIZE: usize = 2;
 
-/// A special version for batch allocation stat, passed in another task `context` C to report task
-/// mem usage 0 bytes at the end.
-#[cfg(enable_task_local_alloc)]
-pub async fn allocation_stat_for_batch<Fut, T, F, C>(
-    future: Fut,
-    interval: Duration,
-    mut report: F,
-    context: C,
-) -> T
-where
-    Fut: Future<Output = T>,
-    F: FnMut(usize),
-    C: BatchTaskContext,
-{
-    use task_stats_alloc::{TaskLocalBytesAllocated, BYTES_ALLOCATED};
-
-    BYTES_ALLOCATED
-        .scope(TaskLocalBytesAllocated::new(), async move {
-            // The guard has the same lifetime as the counter so that the counter will keep positive
-            // in the whole scope. When the scope exits, the guard is released, so the counter can
-            // reach zero eventually and then `drop` itself.
-            let _guard = Box::new(114514);
-            let monitor = async move {
-                let mut interval = tokio::time::interval(interval);
-                loop {
-                    interval.tick().await;
-                    BYTES_ALLOCATED.with(|bytes| report(bytes.val()));
-                }
-            };
-            let output = tokio::select! {
-                biased;
-                _ = monitor => unreachable!(),
-                output = future => {
-                    // NOTE: Report bytes allocated when the actor ends. We simply report 0 here,
-                    // assuming that all memory allocated by this batch task will be freed at some
-                    // time. Maybe we should introduce a better monitoring strategy for batch memory
-                    // usage.
-                    BYTES_ALLOCATED.with(|_| context.store_mem_usage(0));
-                    output
-                },
-            };
-            output
-        })
-        .await
-}
-
 /// Send batch task status (local/distributed) to frontend.
 ///
 ///
@@ -538,27 +488,7 @@ impl<C: BatchTaskContext> BatchTaskExecution<C> {
             }
         };
 
-        #[cfg(enable_task_local_alloc)]
-        {
-            // For every fired Batch Task, we will wrap it with allocation stats to report memory
-            // estimation per task to `BatchManager`.
-            let ctx1 = self.context.clone();
-            let ctx2 = self.context.clone();
-            let alloc_stat_wrap_fut = allocation_stat_for_batch(
-                fut,
-                Duration::from_millis(1000),
-                move |bytes| {
-                    ctx1.store_mem_usage(bytes);
-                },
-                ctx2,
-            );
-            self.runtime.spawn(alloc_stat_wrap_fut);
-        }
-
-        #[cfg(not(enable_task_local_alloc))]
-        {
-            self.runtime.spawn(fut);
-        }
+        self.runtime.spawn(fut);
 
         Ok(())
     }

src/cmd/Cargo.toml

@@ -46,9 +46,6 @@ tokio = { version = "0.2", package = "madsim-tokio", features = [
 workspace-config = { path = "../utils/workspace-config" }
 workspace-hack = { path = "../workspace-hack" }
 
-[target.'cfg(enable_task_local_alloc)'.dependencies]
-task_stats_alloc = { path = "../utils/task_stats_alloc" }
-
 [lib]
 test = false
 

src/cmd/src/lib.rs

@@ -23,10 +23,6 @@ use risingwave_rt::{init_risingwave_logger, main_okk, LoggerSettings};
 #[macro_export]
 macro_rules! main {
     ($component:ident) => {
-        #[cfg(enable_task_local_alloc)]
-        risingwave_common::enable_task_local_jemalloc!();
-
-        #[cfg(not(enable_task_local_alloc))]
         risingwave_common::enable_jemalloc!();
 
         #[cfg_attr(coverage, coverage(off))]

src/cmd_all/Cargo.toml

@@ -67,9 +67,6 @@ vergen = { version = "8", default-features = false, features = [
     "gitcl",
 ] }
 
-[target.'cfg(enable_task_local_alloc)'.dependencies]
-task_stats_alloc = { path = "../utils/task_stats_alloc" }
-
 [[bin]]
 name = "risingwave"
 path = "src/bin/risingwave.rs"

src/cmd_all/src/bin/risingwave.rs

@@ -31,10 +31,6 @@ use strum::IntoEnumIterator;
 use strum_macros::{Display, EnumIter, EnumString, IntoStaticStr};
 use tracing::Level;
 
-#[cfg(enable_task_local_alloc)]
-risingwave_common::enable_task_local_jemalloc!();
-
-#[cfg(not(enable_task_local_alloc))]
 risingwave_common::enable_jemalloc!();
 
 const BINARY_NAME: &str = "risingwave";

src/stream/Cargo.toml

@@ -81,9 +81,6 @@ tokio-stream = "0.1"
 tonic = { workspace = true }
 tracing = "0.1"
 
-[target.'cfg(enable_task_local_alloc)'.dependencies]
-task_stats_alloc = { path = "../utils/task_stats_alloc" }
-
 [target.'cfg(not(madsim))'.dependencies]
 workspace-hack = { path = "../workspace-hack" }
 

src/stream/src/task/stream_manager.rs

@@ -651,30 +651,7 @@ impl LocalBarrierWorker {
                 let with_config =
                     crate::CONFIG.scope(self.actor_manager.env.config().clone(), instrumented);
 
-                // TODO: are the following lines still useful?
-                #[cfg(enable_task_local_alloc)]
-                {
-                    let metrics = streaming_metrics.clone();
-                    let actor_id_str = actor_id.to_string();
-                    let fragment_id_str = actor_context.fragment_id.to_string();
-                    let allocation_stated = task_stats_alloc::allocation_stat(
-                        with_config,
-                        Duration::from_millis(1000),
-                        move |bytes| {
-                            metrics
-                                .actor_memory_usage
-                                .with_label_values(&[&actor_id_str, &fragment_id_str])
-                                .set(bytes as i64);
-
-                            actor_context.store_mem_usage(bytes);
-                        },
-                    );
-                    self.actor_manager.runtime.spawn(allocation_stated)
-                }
-                #[cfg(not(enable_task_local_alloc))]
-                {
-                    self.actor_manager.runtime.spawn(with_config)
-                }
+                self.actor_manager.runtime.spawn(with_config)
             };
             self.actor_manager_state.handles.insert(actor_id, handle);