execution_plans: add MetricsCollectingStream

This change adds a new type called `MetricsCollectingStream`. It wraps a stream of `FlightData`
and collects any metrics that are passed in the `app_metadata`. This change also introduces
an `FlightAppMetadata` enum proto which can be used to define our app metadata protocol.

Author: jayshrivastava

src/execution_plans/metrics_collecting_stream.rs

@@ -0,0 +1,311 @@
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use crate::metrics::proto::MetricsSetProto;
+use crate::protobuf::StageKey;
+use crate::protobuf::{AppMetadata, FlightAppMetadata};
+use arrow_flight::{error::FlightError, FlightData};
+use dashmap::DashMap;
+use futures::stream::Stream;
+use prost::Message;
+use std::sync::Arc;
+
+/// MetricsCollectingStream wraps a FlightData stream and extracts metrics from app_metadata
+/// while passing through all the other FlightData unchanged.
+pub struct MetricsCollectingStream<S>
+where
+    S: Stream<Item = Result<FlightData, FlightError>> + Send + Unpin,
+{
+    inner: S,
+    metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
+}
+
+impl<S> MetricsCollectingStream<S>
+where
+    S: Stream<Item = Result<FlightData, FlightError>> + Send + Unpin,
+{
+    #[allow(dead_code)]
+    pub fn new(
+        stream: S,
+        metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
+    ) -> Self {
+        Self {
+            inner: stream,
+            metrics_collection,
+        }
+    }
+
+    fn extract_metrics_from_flight_data(
+        &self,
+        flight_data: &mut FlightData,
+    ) -> Result<(), FlightError> {
+        if !flight_data.app_metadata.is_empty() {
+            return match FlightAppMetadata::decode(flight_data.app_metadata.as_ref()) {
+                Ok(metadata) => {
+                    if let Some(content) = metadata.content {
+                        match content {
+                            AppMetadata::MetricsCollection(task_metrics_set) => {
+                                for task_metrics in task_metrics_set.tasks {
+                                    if let Some(stage_key) = task_metrics.stage_key {
+                                        self.metrics_collection
+                                            .insert(stage_key, task_metrics.metrics);
+                                    } else {
+                                        return Err(FlightError::ProtocolError("expected Some StageKey in MetricsCollectingStream, got None".to_string()));
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    flight_data.app_metadata.clear();
+                    Ok(())
+                }
+                Err(e) => Err(FlightError::ProtocolError(format!(
+                    "failed to decode app_metadata: {}",
+                    e
+                ))),
+            };
+        }
+        Ok(())
+    }
+}
+
+impl<S> Stream for MetricsCollectingStream<S>
+where
+    S: Stream<Item = Result<FlightData, FlightError>> + Send + Unpin,
+{
+    type Item = Result<FlightData, FlightError>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match Pin::new(&mut self.inner).poll_next(cx) {
+            Poll::Ready(Some(Ok(mut flight_data))) => {
+                // Extract metrics from app_metadata if present.
+                match self.extract_metrics_from_flight_data(&mut flight_data) {
+                    Ok(_) => Poll::Ready(Some(Ok(flight_data))),
+                    Err(e) => Poll::Ready(Some(Err(e))),
+                }
+            }
+            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::metrics::proto::{
+        MetricProto, MetricValueProto, MetricsSetProto, NamedCount, NamedGauge,
+    };
+    use crate::protobuf::{
+        AppMetadata, FlightAppMetadata, MetricsCollection, StageKey, TaskMetrics,
+    };
+    use arrow_flight::FlightData;
+    use futures::stream::{self, StreamExt};
+    use prost::{bytes::Bytes, Message};
+
+    #[tokio::test]
+    async fn test_metrics_collecting_stream_extracts_and_removes_metadata() {
+        let test_metrics_set = MetricsSetProto {
+            metrics: vec![
+                MetricProto {
+                    metric: Some(MetricValueProto::Count(NamedCount {
+                        name: "test_count".to_string(),
+                        value: 42,
+                    })),
+                    labels: vec![],
+                    partition: Some(0),
+                },
+                MetricProto {
+                    metric: Some(MetricValueProto::Gauge(NamedGauge {
+                        name: "test_gauge".to_string(),
+                        value: 99,
+                    })),
+                    labels: vec![],
+                    partition: Some(0),
+                },
+            ],
+        };
+
+        let stage_keys = vec![
+            StageKey {
+                query_id: "test_query".to_string(),
+                stage_id: 1,
+                task_number: 1,
+            },
+            StageKey {
+                query_id: "test_query_2".to_string(),
+                stage_id: 2,
+                task_number: 2,
+            },
+        ];
+
+        let app_metadatas = stage_keys
+            .iter()
+            .map(|stage_key| FlightAppMetadata {
+                content: Some(AppMetadata::MetricsCollection(MetricsCollection {
+                    tasks: vec![TaskMetrics {
+                        stage_key: Some(stage_key.clone()),
+                        metrics: vec![test_metrics_set.clone()],
+                    }],
+                })),
+            })
+            .collect::<Vec<_>>();
+
+        // Create test FlightData messages - some with metadata, some without
+        let flight_data_with_metadata = FlightData {
+            flight_descriptor: None,
+            data_header: Bytes::new(),
+            app_metadata: app_metadatas[0].encode_to_vec().into(),
+            data_body: vec![1, 2, 3].into(),
+        };
+
+        let flight_data_without_metadata = FlightData {
+            flight_descriptor: None,
+            data_header: Bytes::new(),
+            app_metadata: Bytes::new(),
+            data_body: vec![4, 5, 6].into(),
+        };
+
+        let flight_data_with_metadata2 = FlightData {
+            flight_descriptor: None,
+            data_header: Bytes::new(),
+            app_metadata: app_metadatas[1].encode_to_vec().into(),
+            data_body: vec![7, 8, 9].into(),
+        };
+
+        let input_stream = stream::iter(vec![
+            Ok(flight_data_with_metadata),
+            Ok(flight_data_without_metadata),
+            Ok(flight_data_with_metadata2),
+        ]);
+
+        let metrics_collection = Arc::new(DashMap::new());
+        let mut collecting_stream =
+            MetricsCollectingStream::new(input_stream, metrics_collection.clone());
+
+        // Collect all messages from the stream. All should have empty app_metadata.
+        let mut collected_messages = vec![];
+        while let Some(result) = collecting_stream.next().await {
+            collected_messages.push(result.unwrap());
+        }
+        assert_eq!(collected_messages.len(), 3);
+        for msg in &collected_messages {
+            assert!(
+                msg.app_metadata.is_empty(),
+                "app_metadata should be empty after collection"
+            );
+        }
+
+        // Verify the data in the messages.
+        assert_eq!(collected_messages[0].data_body, vec![1, 2, 3]);
+        assert_eq!(collected_messages[1].data_body, vec![4, 5, 6]);
+        assert_eq!(collected_messages[2].data_body, vec![7, 8, 9]);
+
+        // Verify metrics were collected
+        assert_eq!(metrics_collection.len(), 2);
+        for stage_key in stage_keys {
+            let collected_metrics = metrics_collection.get(&stage_key).unwrap();
+            assert_eq!(collected_metrics.len(), 1);
+            assert_eq!(collected_metrics[0].metrics.len(), 2); // We have 2 metrics: Count and Gauge
+                                                               // Verify the first metric value (Count)
+            if let Some(MetricValueProto::Count(count)) = &collected_metrics[0].metrics[0].metric {
+                assert_eq!(count.name, "test_count");
+                assert_eq!(count.value, 42);
+            } else {
+                panic!("expected Count metric");
+            }
+
+            // Verify the second metric value (Gauge)
+            if let Some(MetricValueProto::Gauge(gauge)) = &collected_metrics[0].metrics[1].metric {
+                assert_eq!(gauge.name, "test_gauge");
+                assert_eq!(gauge.value, 99);
+            } else {
+                panic!("expected Gauge metric");
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_metrics_collecting_stream_error_missing_stage_key() {
+        let metrics_collection = Arc::new(DashMap::new());
+        let task_metrics_with_no_stage_key = TaskMetrics {
+            stage_key: None,
+            metrics: vec![MetricsSetProto {
+                metrics: vec![MetricProto {
+                    metric: Some(MetricValueProto::Count(NamedCount {
+                        name: "test_count".to_string(),
+                        value: 42,
+                    })),
+                    labels: vec![],
+                    partition: Some(0),
+                }],
+            }],
+        };
+
+        let invalid_app_metadata = FlightAppMetadata {
+            content: Some(AppMetadata::MetricsCollection(MetricsCollection {
+                tasks: vec![task_metrics_with_no_stage_key],
+            })),
+        };
+
+        let invalid_flight_data = FlightData {
+            flight_descriptor: None,
+            data_header: Bytes::new(),
+            app_metadata: invalid_app_metadata.encode_to_vec().into(),
+            data_body: vec![1, 2, 3].into(),
+        };
+
+        let error_stream = stream::iter(vec![Ok(invalid_flight_data)]);
+        let mut collecting_stream = MetricsCollectingStream::new(error_stream, metrics_collection);
+
+        let result = collecting_stream.next().await.unwrap();
+        assert!(result.is_err());
+        if let Err(FlightError::ProtocolError(msg)) = result {
+            assert!(msg.contains("expected Some StageKey in MetricsCollectingStream, got None"));
+        } else {
+            panic!("expected FlightError::ProtocolError with stage key error");
+        }
+    }
+
+    #[tokio::test]
+    async fn test_metrics_collecting_stream_error_invalid_metadata() {
+        let metrics_collection = Arc::new(DashMap::new());
+
+        let flight_data_with_invalid_metadata = FlightData {
+            flight_descriptor: None,
+            data_header: Bytes::new(),
+            app_metadata: vec![0xFF, 0xFF, 0xFF, 0xFF].into(), // Invalid protobuf data
+            data_body: vec![4, 5, 6].into(),
+        };
+
+        let error_stream = stream::iter(vec![Ok(flight_data_with_invalid_metadata)]);
+        let mut collecting_stream = MetricsCollectingStream::new(error_stream, metrics_collection);
+
+        let result = collecting_stream.next().await.unwrap();
+        assert!(result.is_err());
+        if let Err(FlightError::ProtocolError(msg)) = result {
+            assert!(msg.contains("failed to decode app_metadata"));
+        } else {
+            panic!("expected FlightError::ProtocolError with decode error");
+        }
+    }
+
+    #[tokio::test]
+    async fn test_metrics_collecting_stream_error_propagation() {
+        let metrics_collection = Arc::new(DashMap::new());
+
+        // Create a stream that emits an error - should be propagated through
+        let stream_error = FlightError::ProtocolError("stream error from inner stream".to_string());
+        let error_stream = stream::iter(vec![Err(stream_error)]);
+        let mut collecting_stream = MetricsCollectingStream::new(error_stream, metrics_collection);
+
+        let result = collecting_stream.next().await.unwrap();
+        assert!(result.is_err());
+        if let Err(FlightError::ProtocolError(msg)) = result {
+            assert!(msg.contains("stream error from inner stream"));
+        } else {
+            panic!("expected FlightError::ProtocolError with inner stream error");
+        }
+    }
+}

src/execution_plans/mod.rs

@@ -1,5 +1,6 @@
 mod arrow_flight_read;
 mod metrics;
+mod metrics_collecting_stream;
 mod partition_isolator;
 mod stage;
 

src/protobuf/app_metadata.rs

@@ -0,0 +1,39 @@
+use crate::metrics::proto::MetricsSetProto;
+use crate::protobuf::StageKey;
+
+/// A collection of metrics for a set of tasks in an ExecutionPlan. each
+/// entry should have a distinct [StageKey].
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct MetricsCollection {
+    #[prost(message, repeated, tag = "1")]
+    pub tasks: Vec<TaskMetrics>,
+}
+
+/// TaskMetrics represents the metrics for a single task.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct TaskMetrics {
+    /// stage_key uniquely identifies this task.
+    ///
+    /// This field is always present. It's marked optional due to protobuf rules.
+    #[prost(message, optional, tag = "1")]
+    pub stage_key: Option<StageKey>,
+    /// metrics[i] is the set of metrics for plan node `i` where plan nodes are in pre-order
+    /// traversal order.
+    #[prost(message, repeated, tag = "2")]
+    pub metrics: Vec<MetricsSetProto>,
+}
+
+// FlightAppMetadata represents all types of app_metadata which we use in the distributed execution.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct FlightAppMetadata {
+    #[prost(oneof = "AppMetadata", tags = "1")]
+    pub content: Option<AppMetadata>,
+}
+
+#[derive(Clone, PartialEq, ::prost::Oneof)]
+pub enum AppMetadata {
+    #[prost(message, tag = "1")]
+    MetricsCollection(MetricsCollection),
+    // Note: For every additional enum variant, ensure to add tags to [FlightAppMetadata]. ex. `#[prost(oneof = "AppMetadata", tags = "1,2,3")]` etc.
+    // If you don't the proto will compile but you may encounter errors during serialization/deserialization.
+}

src/protobuf/mod.rs

@@ -1,7 +1,10 @@
+mod app_metadata;
 mod distributed_codec;
 mod stage_proto;
 mod user_codec;
 
+#[allow(unused_imports)]
+pub(crate) use app_metadata::{AppMetadata, FlightAppMetadata, MetricsCollection, TaskMetrics};
 pub(crate) use distributed_codec::DistributedCodec;
 pub(crate) use stage_proto::{proto_from_stage, stage_from_proto, StageExecProto, StageKey};
 pub(crate) use user_codec::{get_distributed_user_codec, set_distributed_user_codec};