refactor: bring metrics to http output (#3247)

* refactor: bring metrics to http output

* chore: remove unwrap

* chore: make walk plan accumulate

* chore: change field name and comment

* chore: add metrics to http resp header

* chore: move PrometheusJsonResponse to a separate file and impl IntoResponse

* chore: put metrics in prometheus resp header too

src/common/plugins/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "common-plugins"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]

src/common/plugins/src/consts.rs

@@ -0,0 +1,19 @@
+// Copyright 2023 Greptime Team
+//
+// 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.
+
+/// Greptime preserved metrics prefix
+pub const GREPTIME_EXEC_PREFIX: &str = "greptime_exec_";
+
+/// Execution cost metrics key
+pub const GREPTIME_EXEC_COST: &str = "greptime_exec_cost";

src/common/plugins/src/lib.rs

@@ -0,0 +1,20 @@
+// Copyright 2023 Greptime Team
+//
+// 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.
+
+/// This crate is designed to be at the bottom of the depencey tree
+/// to provide common and useful utils and consts to all plugin usage,
+/// since `plugins` crate is at the top depending on crates like `frontend` and `datanode`
+mod consts;
+
+pub use consts::{GREPTIME_EXEC_COST, GREPTIME_EXEC_PREFIX};

src/common/query/src/lib.rs

@@ -13,10 +13,12 @@
 // limitations under the License.
 
 use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
 
 use api::greptime_proto::v1::add_column_location::LocationType;
 use api::greptime_proto::v1::AddColumnLocation as Location;
 use common_recordbatch::{RecordBatches, SendableRecordBatchStream};
+use physical_plan::PhysicalPlan;
 use serde::{Deserialize, Serialize};
 
 pub mod columnar_value;
@@ -32,7 +34,14 @@ use sqlparser_derive::{Visit, VisitMut};
 pub enum Output {
     AffectedRows(usize),
     RecordBatches(RecordBatches),
-    Stream(SendableRecordBatchStream),
+    Stream(SendableRecordBatchStream, Option<Arc<dyn PhysicalPlan>>),
+}
+
+impl Output {
+    // helper function to build original `Output::Stream`
+    pub fn new_stream(stream: SendableRecordBatchStream) -> Self {
+        Output::Stream(stream, None)
+    }
 }
 
 impl Debug for Output {
@@ -42,7 +51,13 @@ impl Debug for Output {
             Output::RecordBatches(recordbatches) => {
                 write!(f, "Output::RecordBatches({recordbatches:?})")
             }
-            Output::Stream(_) => write!(f, "Output::Stream(<stream>)"),
+            Output::Stream(_, df) => {
+                if df.is_some() {
+                    write!(f, "Output::Stream(<stream>, Some<physical_plan>)")
+                } else {
+                    write!(f, "Output::Stream(<stream>)")
+                }
+            }
         }
     }
 }

src/common/recordbatch/src/adapter.rs

@@ -182,7 +182,7 @@ pub struct RecordBatchStreamAdapter {
 enum Metrics {
     Unavailable,
     Unresolved(Arc<dyn ExecutionPlan>),
-    Resolved(String),
+    Resolved(RecordBatchMetrics),
 }
 
 impl RecordBatchStreamAdapter {
@@ -222,9 +222,9 @@ impl RecordBatchStream for RecordBatchStreamAdapter {
         self.schema.clone()
     }
 
-    fn metrics(&self) -> Option<String> {
+    fn metrics(&self) -> Option<RecordBatchMetrics> {
         match &self.metrics_2 {
-            Metrics::Resolved(metrics) => Some(metrics.clone()),
+            Metrics::Resolved(metrics) => Some(*metrics),
             Metrics::Unavailable | Metrics::Unresolved(_) => None,
         }
     }
@@ -254,8 +254,7 @@ impl Stream for RecordBatchStreamAdapter {
                     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());
+                        self.metrics_2 = Metrics::Resolved(metrics_holder);
                     }
                 }
                 Poll::Ready(None)
@@ -285,7 +284,7 @@ fn collect_metrics(df_plan: &Arc<dyn ExecutionPlan>, result: &mut RecordBatchMet
 
 /// [`RecordBatchMetrics`] carrys metrics value
 /// from datanode to frontend through gRPC
-#[derive(serde::Serialize, serde::Deserialize, Default, Debug)]
+#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone, Copy)]
 pub struct RecordBatchMetrics {
     // cpu consumption in nanoseconds
     pub elapsed_compute: usize,

src/common/recordbatch/src/lib.rs

@@ -21,6 +21,7 @@ pub mod util;
 use std::pin::Pin;
 use std::sync::Arc;
 
+use adapter::RecordBatchMetrics;
 use arc_swap::ArcSwapOption;
 use datafusion::physical_plan::memory::MemoryStream;
 pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
@@ -42,7 +43,7 @@ pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
         None
     }
 
-    fn metrics(&self) -> Option<String> {
+    fn metrics(&self) -> Option<RecordBatchMetrics> {
         None
     }
 }
@@ -212,7 +213,7 @@ pub struct RecordBatchStreamWrapper<S> {
     pub schema: SchemaRef,
     pub stream: S,
     pub output_ordering: Option<Vec<OrderOption>>,
-    pub metrics: Arc<ArcSwapOption<String>>,
+    pub metrics: Arc<ArcSwapOption<RecordBatchMetrics>>,
 }
 
 impl<S> RecordBatchStreamWrapper<S> {
@@ -238,8 +239,8 @@ impl<S: Stream<Item = Result<RecordBatch>> + Unpin> RecordBatchStream
         self.output_ordering.as_deref()
     }
 
-    fn metrics(&self) -> Option<String> {
-        self.metrics.load().as_ref().map(|s| s.as_ref().clone())
+    fn metrics(&self) -> Option<RecordBatchMetrics> {
+        self.metrics.load().as_ref().map(|s| *s.as_ref())
     }
 }
 

src/common/test-util/src/recordbatch.rs

@@ -28,7 +28,7 @@ pub async fn execute_and_check_output(db: &Database, sql: &str, expected: Expect
             assert_eq!(*x, y, "actual: \n{}", x)
         }
         (Output::RecordBatches(_), ExpectedOutput::QueryResult(x))
-        | (Output::Stream(_), ExpectedOutput::QueryResult(x)) => {
+        | (Output::Stream(_, _), ExpectedOutput::QueryResult(x)) => {
             check_output_stream(output, x).await
         }
         _ => panic!(),
@@ -37,7 +37,7 @@ pub async fn execute_and_check_output(db: &Database, sql: &str, expected: Expect
 
 pub async fn check_output_stream(output: Output, expected: &str) {
     let recordbatches = match output {
-        Output::Stream(stream) => util::collect_batches(stream).await.unwrap(),
+        Output::Stream(stream, _) => util::collect_batches(stream).await.unwrap(),
         Output::RecordBatches(recordbatches) => recordbatches,
         _ => unreachable!(),
     };