chore: carry metrics in flight metadata from datanode to frontend (#3113)

* chore: carry metrics in flight metadata from datanode to frontend

* chore: fix typo

* fix: ignore metric flight message on client

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* chore: add cr comment

Co-authored-by: Ruihang Xia <waynestxia@gmail.com>

* chore: add cr comment

Co-authored-by: Ruihang Xia <waynestxia@gmail.com>

* chore: update proto

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Co-authored-by: Ruihang Xia <waynestxia@gmail.com>

src/common/recordbatch/Cargo.toml

@@ -5,6 +5,7 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
+arc-swap = "1.6"
 common-error.workspace = true
 common-macro.workspace = true
 datafusion-common.workspace = true
@@ -13,8 +14,8 @@ datatypes.workspace = true
 futures.workspace = true
 paste = "1.0"
 serde.workspace = true
+serde_json.workspace = true
 snafu.workspace = true
 
 [dev-dependencies]
-serde_json = "1.0"
 tokio.workspace = true

src/common/recordbatch/src/adapter.rs

@@ -21,8 +21,8 @@ use datafusion::arrow::compute::cast;
 use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
 use datafusion::error::Result as DfResult;
 use datafusion::parquet::arrow::async_reader::{AsyncFileReader, ParquetRecordBatchStream};
-use datafusion::physical_plan::metrics::BaselineMetrics;
-use datafusion::physical_plan::RecordBatchStream as DfRecordBatchStream;
+use datafusion::physical_plan::metrics::{BaselineMetrics, MetricValue};
+use datafusion::physical_plan::{ExecutionPlan, RecordBatchStream as DfRecordBatchStream};
 use datafusion_common::DataFusionError;
 use datatypes::schema::{Schema, SchemaRef};
 use futures::ready;
@@ -154,6 +154,15 @@ pub struct RecordBatchStreamAdapter {
     schema: SchemaRef,
     stream: DfSendableRecordBatchStream,
     metrics: Option<BaselineMetrics>,
+    /// Aggregated plan-level metrics. Resolved after an [ExecutionPlan] is finished.
+    metrics_2: Metrics,
+}
+
+/// Json encoded metrics. Contains metric from a whole plan tree.
+enum Metrics {
+    Unavailable,
+    Unresolved(Arc<dyn ExecutionPlan>),
+    Resolved(String),
 }
 
 impl RecordBatchStreamAdapter {
@@ -164,12 +173,14 @@ impl RecordBatchStreamAdapter {
             schema,
             stream,
             metrics: None,
+            metrics_2: Metrics::Unavailable,
         })
     }
 
-    pub fn try_new_with_metrics(
+    pub fn try_new_with_metrics_and_df_plan(
         stream: DfSendableRecordBatchStream,
         metrics: BaselineMetrics,
+        df_plan: Arc<dyn ExecutionPlan>,
     ) -> Result<Self> {
         let schema =
             Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
@@ -177,14 +188,26 @@ impl RecordBatchStreamAdapter {
             schema,
             stream,
             metrics: Some(metrics),
+            metrics_2: Metrics::Unresolved(df_plan),
         })
     }
+
+    pub fn set_metrics2(&mut self, plan: Arc<dyn ExecutionPlan>) {
+        self.metrics_2 = Metrics::Unresolved(plan)
+    }
 }
 
 impl RecordBatchStream for RecordBatchStreamAdapter {
     fn schema(&self) -> SchemaRef {
         self.schema.clone()
     }
+
+    fn metrics(&self) -> Option<String> {
+        match &self.metrics_2 {
+            Metrics::Resolved(metrics) => Some(metrics.clone()),
+            Metrics::Unavailable | Metrics::Unresolved(_) => None,
+        }
+    }
 }
 
 impl Stream for RecordBatchStreamAdapter {
@@ -206,7 +229,17 @@ impl Stream for RecordBatchStreamAdapter {
                     df_record_batch,
                 )))
             }
-            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Ready(None) => {
+                if let Metrics::Unresolved(df_plan) = &self.metrics_2 {
+                    let mut metrics_holder = RecordBatchMetrics::default();
+                    collect_metrics(df_plan, &mut metrics_holder);
+                    if metrics_holder.elapsed_compute != 0 || metrics_holder.memory_usage != 0 {
+                        self.metrics_2 =
+                            Metrics::Resolved(serde_json::to_string(&metrics_holder).unwrap());
+                    }
+                }
+                Poll::Ready(None)
+            }
         }
     }
 
@@ -216,6 +249,30 @@ impl Stream for RecordBatchStreamAdapter {
     }
 }
 
+fn collect_metrics(df_plan: &Arc<dyn ExecutionPlan>, result: &mut RecordBatchMetrics) {
+    if let Some(metrics) = df_plan.metrics() {
+        metrics.iter().for_each(|m| match m.value() {
+            MetricValue::ElapsedCompute(ec) => result.elapsed_compute += ec.value(),
+            MetricValue::CurrentMemoryUsage(m) => result.memory_usage += m.value(),
+            _ => {}
+        });
+    }
+
+    for child in df_plan.children() {
+        collect_metrics(&child, result);
+    }
+}
+
+/// [`RecordBatchMetrics`] carrys metrics value
+/// from datanode to frontend through gRPC
+#[derive(serde::Serialize, serde::Deserialize, Default, Debug)]
+pub struct RecordBatchMetrics {
+    // cpu consumption in nanoseconds
+    pub elapsed_compute: usize,
+    // memory used by the plan in bytes
+    pub memory_usage: usize,
+}
+
 enum AsyncRecordBatchStreamAdapterState {
     Uninit(FutureStream),
     Ready(SendableRecordBatchStream),

src/common/recordbatch/src/lib.rs

@@ -20,6 +20,7 @@ pub mod util;
 use std::pin::Pin;
 use std::sync::Arc;
 
+use arc_swap::ArcSwapOption;
 use datafusion::physical_plan::memory::MemoryStream;
 pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
 use datatypes::arrow::compute::SortOptions;
@@ -39,6 +40,10 @@ pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
     fn output_ordering(&self) -> Option<&[OrderOption]> {
         None
     }
+
+    fn metrics(&self) -> Option<String> {
+        None
+    }
 }
 
 pub type SendableRecordBatchStream = Pin<Box<dyn RecordBatchStream + Send>>;
@@ -206,6 +211,7 @@ pub struct RecordBatchStreamWrapper<S> {
     pub schema: SchemaRef,
     pub stream: S,
     pub output_ordering: Option<Vec<OrderOption>>,
+    pub metrics: Arc<ArcSwapOption<String>>,
 }
 
 impl<S> RecordBatchStreamWrapper<S> {
@@ -215,6 +221,7 @@ impl<S> RecordBatchStreamWrapper<S> {
             schema,
             stream,
             output_ordering: None,
+            metrics: Default::default(),
         }
     }
 }
@@ -229,6 +236,10 @@ impl<S: Stream<Item = Result<RecordBatch>> + Unpin> RecordBatchStream
     fn output_ordering(&self) -> Option<&[OrderOption]> {
         self.output_ordering.as_deref()
     }
+
+    fn metrics(&self) -> Option<String> {
+        self.metrics.load().as_ref().map(|s| s.as_ref().clone())
+    }
 }
 
 impl<S: Stream<Item = Result<RecordBatch>> + Unpin> Stream for RecordBatchStreamWrapper<S> {