execution_plans: add metrics collector and re-writer

This change introduces new structs / concepts:

1. MetricsCollector

Collects metrics from an instance of `StageExec` (which is a task) using a
pre-order traversal. It stops at any `ArrowFlightReadExec` nodes and collects
child task metrics from them, if there are any.

2. MetricsWrapperExec

A new `ExecutionPlan` node which cannot be executed. It wraps a "real" `ExecutionPlan`
node and stores metrics. This let's you "override" the `metrics()` method on
`ExecutionPlan` nodes. This override applies when displaying the plan as well.

This struct is private.

3. Task MetricsRewriter

Rewrites a task by wrapping each node in a `MetricsWrapperExec`.

Informs #123.

Author: jayshrivastava

src/execution_plans/arrow_flight_read.rs

@@ -3,12 +3,14 @@ use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error};
 use crate::execution_plans::StageExec;
 use crate::flight_service::DoGet;
+use crate::metrics::proto::MetricsSetProto;
 use crate::protobuf::{proto_from_stage, DistributedCodec, StageKey};
 use crate::ChannelResolver;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::Ticket;
+use dashmap::DashMap;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
@@ -56,6 +58,13 @@ pub struct ArrowFlightReadReadyExec {
     /// the properties we advertise for this execution plan
     properties: PlanProperties,
     pub(crate) stage_num: usize,
+    /// metrics_collection is used to collect metrics from child tasks. It is empty when an ArrowFlightReadReadyExec is instansiated
+    /// (deserialized, created via [ArrowFlightReadExec::new_ready] etc). Metrics are populated in this map via [ArrowFlightReadExec::execute].
+    ///
+    /// An instance may recieve metrics for 0 to N child tasks, where N is the number of tasks in the stage it is reading from.
+    /// This is because, by convention, the ArrowFlightEndpoint sends metrics for a task to the last ArrowFlightReadExec to read from it, which
+    /// may or may not be this instance.
+    pub(super) metrics_collection: Arc<DashMap<StageKey, Vec<MetricsSetProto>>>,
 }
 
 impl ArrowFlightReadExec {
@@ -85,6 +94,7 @@ impl ArrowFlightReadExec {
         Self::Ready(ArrowFlightReadReadyExec {
             properties,
             stage_num,
+            metrics_collection: Arc::new(DashMap::new()),
         })
     }