chore(cubestore): Memory usage tracing (#6663)

Author: waralexrom

rust/cubestore/cubestore/src/cluster/worker_pool.rs

@@ -8,6 +8,7 @@ use std::time::Duration;
 use async_trait::async_trait;
 use deadqueue::unlimited;
 use futures::future::join_all;
+use futures::future::BoxFuture;
 use ipc_channel::ipc;
 use ipc_channel::ipc::{IpcReceiver, IpcSender};
 use log::error;
@@ -334,10 +335,11 @@ where
     let runtime = tokio_builder.build().unwrap();
     worker_setup(&runtime);
     runtime.block_on(async move {
-        let config = Config::default();
-        config.configure_injector().await;
+        let config = get_worker_config().await;
         let services = config.worker_services().await;
 
+        spawn_background_processes(config.clone());
+
         loop {
             let res = rx.recv();
             match res {
@@ -369,15 +371,61 @@ fn worker_setup(runtime: &Runtime) {
     }
 }
 
+async fn get_worker_config() -> Config {
+    let custom_fn = SELECT_WORKER_CONFIGURE_FN.read().unwrap();
+    if let Some(func) = custom_fn.as_ref() {
+        func().await
+    } else {
+        let config = Config::default();
+        config.configure_injector().await;
+        config
+    }
+}
+
+fn spawn_background_processes(config: Config) {
+    let custom_fn = SELECT_WORKER_SPAWN_BACKGROUND_FN.read().unwrap();
+    if let Some(func) = custom_fn.as_ref() {
+        func(config);
+    }
+}
+
 lazy_static! {
     static ref SELECT_WORKER_SETUP: std::sync::RwLock<Option<Box<dyn Fn(&Runtime) + Send + Sync>>> =
         std::sync::RwLock::new(None);
 }
 
+lazy_static! {
+    static ref SELECT_WORKER_CONFIGURE_FN: std::sync::RwLock<Option<Box<dyn Fn() -> BoxFuture<'static, Config> + Send + Sync>>> =
+        std::sync::RwLock::new(None);
+}
+
+lazy_static! {
+    static ref SELECT_WORKER_SPAWN_BACKGROUND_FN: std::sync::RwLock<Option<Box<dyn Fn(Config) + Send + Sync>>> =
+        std::sync::RwLock::new(None);
+}
+
 pub fn register_select_worker_setup(f: fn(&Runtime)) {
-    let mut startup = SELECT_WORKER_SETUP.write().unwrap();
-    assert!(startup.is_none(), "select worker setup already registered");
-    *startup = Some(Box::new(f));
+    let mut setup = SELECT_WORKER_SETUP.write().unwrap();
+    assert!(setup.is_none(), "select worker setup already registered");
+    *setup = Some(Box::new(f));
+}
+
+pub fn register_select_worker_configure_fn(f: fn() -> BoxFuture<'static, Config>) {
+    let mut func = SELECT_WORKER_CONFIGURE_FN.write().unwrap();
+    assert!(
+        func.is_none(),
+        "select worker configure function already registered"
+    );
+    *func = Some(Box::new(f));
+}
+
+pub fn register_select_worker_spawn_background_fn(f: fn(Config)) {
+    let mut func = SELECT_WORKER_SPAWN_BACKGROUND_FN.write().unwrap();
+    assert!(
+        func.is_none(),
+        "select worker spawn background function already registered"
+    );
+    *func = Some(Box::new(f));
 }
 
 #[cfg(test)]

rust/cubestore/cubestore/src/config/mod.rs

@@ -37,6 +37,7 @@ use crate::telemetry::tracing::{TracingHelper, TracingHelperImpl};
 use crate::telemetry::{
     start_agent_event_loop, start_track_event_loop, stop_agent_event_loop, stop_track_event_loop,
 };
+use crate::util::memory::{MemoryHandler, MemoryHandlerImpl};
 use crate::CubeError;
 use datafusion::cube_ext;
 use datafusion::physical_plan::parquet::{LruParquetMetadataCache, NoopParquetMetadataCache};
@@ -1761,7 +1762,7 @@ impl Config {
 
         self.injector
             .register_typed_with_default::<dyn QueryExecutor, _, _, _>(async move |i| {
-                QueryExecutorImpl::new(i.get_service_typed().await)
+                QueryExecutorImpl::new(i.get_service_typed().await, i.get_service_typed().await)
             })
             .await;
 
@@ -1781,6 +1782,10 @@ impl Config {
             })
             .await;
 
+        self.injector
+            .register_typed::<dyn MemoryHandler, _, _, _>(async move |_| MemoryHandlerImpl::new())
+            .await;
+
         let query_cache_to_move = query_cache.clone();
         self.injector
             .register_typed::<dyn SqlService, _, _, _>(async move |i| {

rust/cubestore/cubestore/src/queryplanner/check_memory.rs

@@ -0,0 +1,119 @@
+use crate::util::memory::MemoryHandler;
+use arrow::datatypes::SchemaRef;
+use arrow::error::Result as ArrowResult;
+use arrow::record_batch::RecordBatch;
+use async_trait::async_trait;
+use datafusion::error::DataFusionError;
+use datafusion::physical_plan::{
+    ExecutionPlan, OptimizerHints, Partitioning, RecordBatchStream, SendableRecordBatchStream,
+};
+use flatbuffers::bitflags::_core::any::Any;
+use futures::stream::Stream;
+use futures::StreamExt;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+#[derive(Debug)]
+pub struct CheckMemoryExec {
+    pub input: Arc<dyn ExecutionPlan>,
+    pub memory_handler: Arc<dyn MemoryHandler>,
+}
+
+impl CheckMemoryExec {
+    pub fn new(input: Arc<dyn ExecutionPlan>, memory_handler: Arc<dyn MemoryHandler>) -> Self {
+        Self {
+            input,
+            memory_handler,
+        }
+    }
+}
+
+#[async_trait]
+impl ExecutionPlan for CheckMemoryExec {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        self.input.schema()
+    }
+
+    fn output_partitioning(&self) -> Partitioning {
+        self.input.output_partitioning()
+    }
+
+    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
+        vec![self.input.clone()]
+    }
+
+    fn with_new_children(
+        &self,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        assert_eq!(children.len(), 1);
+        Ok(Arc::new(CheckMemoryExec {
+            input: children.into_iter().next().unwrap(),
+            memory_handler: self.memory_handler.clone(),
+        }))
+    }
+
+    fn output_hints(&self) -> OptimizerHints {
+        self.input.output_hints()
+    }
+
+    async fn execute(
+        &self,
+        partition: usize,
+    ) -> Result<SendableRecordBatchStream, DataFusionError> {
+        if partition >= self.input.output_partitioning().partition_count() {
+            return Err(DataFusionError::Internal(format!(
+                "ExecutionPlanExec invalid partition {}",
+                partition
+            )));
+        }
+
+        let input = self.input.execute(partition).await?;
+        Ok(Box::pin(CheckMemoryStream {
+            schema: self.schema(),
+            memory_handler: self.memory_handler.clone(),
+            input,
+        }))
+    }
+}
+
+struct CheckMemoryStream {
+    schema: SchemaRef,
+    memory_handler: Arc<dyn MemoryHandler>,
+    input: SendableRecordBatchStream,
+}
+
+impl Stream for CheckMemoryStream {
+    type Item = ArrowResult<RecordBatch>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.input.poll_next_unpin(cx).map(|x| match x {
+            Some(Ok(batch)) => {
+                let r = self
+                    .memory_handler
+                    .check_memory()
+                    .map(|_| batch)
+                    .map_err(|e| e.into());
+                Some(r)
+            }
+            other => other,
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // same number of record batches
+        self.input.size_hint()
+    }
+}
+
+impl RecordBatchStream for CheckMemoryStream {
+    /// Get the schema
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+}

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -4,6 +4,7 @@ pub mod panic;
 mod partition_filter;
 mod planning;
 pub use planning::PlanningMeta;
+mod check_memory;
 pub mod pretty_printers;
 pub mod query_executor;
 pub mod serialized_plan;

rust/cubestore/cubestore/src/queryplanner/optimizations/check_memory.rs

@@ -0,0 +1,22 @@
+use crate::queryplanner::check_memory::CheckMemoryExec;
+use crate::queryplanner::query_executor::ClusterSendExec;
+use crate::util::memory::MemoryHandler;
+use datafusion::error::DataFusionError;
+use datafusion::physical_plan::memory::MemoryExec;
+use datafusion::physical_plan::parquet::ParquetExec;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+/// Add `CheckMemoryExec` behind some nodes.
+pub fn add_check_memory_exec(
+    p: Arc<dyn ExecutionPlan>,
+    mem_handler: Arc<dyn MemoryHandler>,
+) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+    let p_any = p.as_any();
+    if p_any.is::<ParquetExec>() || p_any.is::<MemoryExec>() || p_any.is::<ClusterSendExec>() {
+        let memory_check = Arc::new(CheckMemoryExec::new(p, mem_handler.clone()));
+        Ok(memory_check)
+    } else {
+        Ok(p)
+    }
+}

rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs

@@ -1,10 +1,17 @@
+mod check_memory;
+mod distributed_partial_aggregate;
+mod prefer_inplace_aggregates;
+pub mod rewrite_plan;
+
 use crate::cluster::Cluster;
 use crate::queryplanner::optimizations::distributed_partial_aggregate::{
     add_limit_to_workers, push_aggregate_to_workers,
 };
 use crate::queryplanner::optimizations::prefer_inplace_aggregates::try_switch_to_inplace_aggregates;
 use crate::queryplanner::planning::CubeExtensionPlanner;
 use crate::queryplanner::serialized_plan::SerializedPlan;
+use crate::util::memory::MemoryHandler;
+use check_memory::add_check_memory_exec;
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::{ExecutionContextState, QueryPlanner};
 use datafusion::logical_plan::LogicalPlan;
@@ -13,30 +20,33 @@ use datafusion::physical_plan::{ExecutionPlan, PhysicalPlanner};
 use rewrite_plan::rewrite_physical_plan;
 use std::sync::Arc;
 
-mod distributed_partial_aggregate;
-mod prefer_inplace_aggregates;
-pub mod rewrite_plan;
-
 pub struct CubeQueryPlanner {
     cluster: Option<Arc<dyn Cluster>>,
     serialized_plan: Arc<SerializedPlan>,
+    memory_handler: Arc<dyn MemoryHandler>,
 }
 
 impl CubeQueryPlanner {
     pub fn new_on_router(
         cluster: Arc<dyn Cluster>,
         serialized_plan: Arc<SerializedPlan>,
+        memory_handler: Arc<dyn MemoryHandler>,
     ) -> CubeQueryPlanner {
         CubeQueryPlanner {
             cluster: Some(cluster),
             serialized_plan,
+            memory_handler,
         }
     }
 
-    pub fn new_on_worker(serialized_plan: Arc<SerializedPlan>) -> CubeQueryPlanner {
+    pub fn new_on_worker(
+        serialized_plan: Arc<SerializedPlan>,
+        memory_handler: Arc<dyn MemoryHandler>,
+    ) -> CubeQueryPlanner {
         CubeQueryPlanner {
             serialized_plan,
             cluster: None,
+            memory_handler,
         }
     }
 }
@@ -54,14 +64,18 @@ impl QueryPlanner for CubeQueryPlanner {
             })])
             .create_physical_plan(logical_plan, ctx_state)?;
         // TODO: assert there is only a single ClusterSendExec in the plan.
-        finalize_physical_plan(p)
+        finalize_physical_plan(p, self.memory_handler.clone())
     }
 }
 
 fn finalize_physical_plan(
     p: Arc<dyn ExecutionPlan>,
+    memory_handler: Arc<dyn MemoryHandler>,
 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
     let p = rewrite_physical_plan(p.as_ref(), &mut |p| try_switch_to_inplace_aggregates(p))?;
     let p = rewrite_physical_plan(p.as_ref(), &mut |p| push_aggregate_to_workers(p))?;
+    let p = rewrite_physical_plan(p.as_ref(), &mut |p| {
+        add_check_memory_exec(p, memory_handler.clone())
+    })?;
     rewrite_physical_plan(p.as_ref(), &mut |p| add_limit_to_workers(p))
 }

rust/cubestore/cubestore/src/queryplanner/pretty_printers.rs

@@ -19,6 +19,7 @@ use datafusion::physical_plan::sort::SortExec;
 use datafusion::physical_plan::ExecutionPlan;
 use itertools::{repeat_n, Itertools};
 
+use crate::queryplanner::check_memory::CheckMemoryExec;
 use crate::queryplanner::filter_by_key_range::FilterByKeyRangeExec;
 use crate::queryplanner::panic::{PanicWorkerExec, PanicWorkerNode};
 use crate::queryplanner::planning::{ClusterSendNode, Snapshot, WorkerExec};
@@ -50,6 +51,7 @@ pub struct PPOptions {
     pub show_aggregations: bool,
     // Applies only to physical plan.
     pub show_output_hints: bool,
+    pub show_check_memory_nodes: bool,
 }
 
 pub fn pp_phys_plan(p: &dyn ExecutionPlan) -> String {
@@ -279,6 +281,13 @@ fn pp_sort_columns(first_agg: usize, cs: &[SortColumn]) -> String {
 }
 
 fn pp_phys_plan_indented(p: &dyn ExecutionPlan, indent: usize, o: &PPOptions, out: &mut String) {
+    if p.as_any().is::<CheckMemoryExec>() && !o.show_check_memory_nodes {
+        //We don't show CheckMemoryExec in plan by default
+        if let Some(child) = p.children().first() {
+            pp_phys_plan_indented(child.as_ref(), indent, o, out)
+        }
+        return;
+    }
     pp_instance(p, indent, o, out);
     if p.as_any().is::<ClusterSendExec>() {
         // Do not show children of ClusterSend. This is a hack to avoid rewriting all tests.