Merge branch 'main' into gabrielmusat/refactor-distributed-planning

# Conflicts:
#	src/execution_plans/distributed.rs
#	src/execution_plans/network_broadcast.rs
#	src/execution_plans/network_coalesce.rs
#	src/execution_plans/network_shuffle.rs
#	src/metrics/task_metrics_rewriter.rs

Author: gabotechs

benchmarks/cdk/bin/worker.rs

@@ -183,6 +183,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
                             physical,
                             DistributedMetricsFormat::PerTask,
                         )
+                        .await
                         .map_err(err)?;
                         let plan = display_plan_ascii(physical.as_ref(), true);
                         drop(task);

console/examples/tpcds_runner.rs

@@ -189,7 +189,8 @@ async fn run_single_query(
     let batches = stream.try_collect::<Vec<_>>().await?;
     if explain_analyze {
         let output =
-            datafusion_distributed::explain_analyze(plan, DistributedMetricsFormat::Aggregated)?;
+            datafusion_distributed::explain_analyze(plan, DistributedMetricsFormat::Aggregated)
+                .await?;
         println!("{output}");
     }
     Ok(batches)

src/execution_plans/benchmarks/shuffle_bench.rs

@@ -229,7 +229,6 @@ impl ShuffleFixture {
                 )),
                 input_stage: input_stage.clone(),
                 worker_connections: WorkerConnectionPool::new(self.bench.producer_tasks),
-                metrics_collection: Arc::new(Default::default()),
             };
             let task_ctx = Arc::new(task_ctx_with_extension(
                 &self.task_ctx,

src/execution_plans/benchmarks/transport_bench.rs

@@ -281,7 +281,6 @@ impl TransportFixture {
                 worker_connections: crate::worker::WorkerConnectionPool::new(
                     self.bench.producer_tasks,
                 ),
-                metrics_collection: Arc::new(Default::default()),
             };
             let task_ctx = Arc::new(task_ctx_with_extension(
                 &self.task_ctx,

src/execution_plans/distributed.rs

@@ -6,17 +6,20 @@ use crate::networking::get_distributed_worker_resolver;
 use crate::passthrough_headers::get_passthrough_headers;
 use crate::protobuf::{DistributedCodec, tonic_status_to_datafusion_error};
 use crate::stage::{LocalStage, RemoteStage, Stage};
+use crate::worker::generated::worker as pb;
 use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration;
 use crate::worker::generated::worker::{
-    CoordinatorToWorkerMsg, SetPlanRequest, TaskKey, WorkUnit, coordinator_to_worker_msg::Inner,
+    CoordinatorToWorkerMsg, SetPlanRequest, TaskKey, WorkUnit, WorkerToCoordinatorMsg,
+    coordinator_to_worker_msg::Inner, worker_to_coordinator_msg,
 };
 use crate::{
     DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedConfig, DistributedTaskContext,
     DistributedWorkUnitFeedContext, WorkerResolver, get_distributed_channel_resolver,
 };
+use datafusion::common::HashMap;
 use datafusion::common::instant::Instant;
 use datafusion::common::runtime::JoinSet;
-use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
 use datafusion::common::{Result, exec_err};
 use datafusion::common::{exec_datafusion_err, internal_datafusion_err};
 use datafusion::error::DataFusionError;
@@ -41,12 +44,46 @@ use std::sync::Mutex;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
 use tokio::sync::mpsc::UnboundedSender;
+use tokio::sync::watch;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tonic::Request;
 use tonic::metadata::MetadataMap;
 use url::Url;
 use uuid::Uuid;
 
+/// Stores the metrics collected from all worker tasks, and notifies waiters when new entries arrive.
+#[derive(Debug, Clone)]
+pub struct MetricsStore {
+    tx: watch::Sender<HashMap<TaskKey, Vec<pb::MetricsSet>>>,
+    rx: watch::Receiver<HashMap<TaskKey, Vec<pb::MetricsSet>>>,
+}
+
+impl MetricsStore {
+    fn new() -> Self {
+        let (tx, rx) = watch::channel(HashMap::new());
+        Self { tx, rx }
+    }
+
+    pub fn insert(&self, key: TaskKey, metrics: Vec<pb::MetricsSet>) {
+        self.tx.send_modify(|map| {
+            map.insert(key, metrics);
+        });
+    }
+
+    pub fn get(&self, key: &TaskKey) -> Option<Vec<pb::MetricsSet>> {
+        self.rx.borrow().get(key).cloned()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn from_entries(
+        entries: impl IntoIterator<Item = (TaskKey, Vec<pb::MetricsSet>)>,
+    ) -> Self {
+        let map: HashMap<_, _> = entries.into_iter().collect();
+        let (tx, rx) = watch::channel(map);
+        Self { tx, rx }
+    }
+}
+
 /// [ExecutionPlan] that executes the inner plan in distributed mode.
 /// Before executing it, two modifications are lazily performed on the plan:
 /// 1. Assigns worker URLs to all the stages. A random set of URLs are sampled from the
@@ -58,6 +95,7 @@ pub struct DistributedExec {
     pub plan: Arc<dyn ExecutionPlan>,
     pub prepared_plan: Arc<Mutex<Option<Arc<dyn ExecutionPlan>>>>,
     metrics: ExecutionPlanMetricsSet,
+    pub task_metrics: Arc<MetricsStore>,
 }
 
 struct PreparedPlan {
@@ -71,9 +109,41 @@ impl DistributedExec {
             plan,
             prepared_plan: Arc::new(Mutex::new(None)),
             metrics: ExecutionPlanMetricsSet::new(),
+            task_metrics: Arc::new(MetricsStore::new()),
         }
     }
 
+    /// Waits until all worker tasks have reported their metrics back via the coordinator channel.
+    ///
+    /// Metrics are delivered asynchronously after query execution completes, so callers that need
+    /// complete metrics (e.g. for observability or display) should await this before inspecting
+    /// [`Self::task_metrics`] or calling [`rewrite_distributed_plan_with_metrics`].
+    ///
+    /// [`rewrite_distributed_plan_with_metrics`]: crate::rewrite_distributed_plan_with_metrics
+    pub async fn wait_for_metrics(&self) {
+        let mut expected_keys: Vec<TaskKey> = Vec::new();
+        let _ = self.plan.apply(|plan| {
+            if let Some(boundary) = plan.as_network_boundary() {
+                let stage = boundary.input_stage();
+                for i in 0..stage.task_count() {
+                    expected_keys.push(TaskKey {
+                        query_id: serialize_uuid(&stage.query_id()),
+                        stage_id: stage.num() as u64,
+                        task_number: i as u64,
+                    });
+                }
+            }
+            Ok(TreeNodeRecursion::Continue)
+        });
+        if expected_keys.is_empty() {
+            return;
+        }
+        let mut rx = self.task_metrics.rx.clone();
+        let _ = rx
+            .wait_for(|map| expected_keys.iter().all(|key| map.contains_key(key)))
+            .await;
+    }
+
     /// Returns the plan which is lazily prepared on execute() and actually gets executed.
     /// It is updated on every call to execute(). Returns an error if .execute() has not been called.
     pub(crate) fn prepared_plan(&self) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
@@ -93,6 +163,8 @@ impl DistributedExec {
     ///    network call feeding the subplan is necessary.
     /// 3. In each network boundary, set the input plan to `None`. That way, network boundaries
     ///    become nodes without children and traversing them will not go further down in.
+    /// 4. Spawn a background task per worker that waits for the worker to finish and collects
+    ///    its metrics into [DistributedExec::task_metrics] via the coordinator channel.
     fn prepare_plan(&self, ctx: &Arc<TaskContext>) -> Result<PreparedPlan> {
         let worker_resolver = get_distributed_worker_resolver(ctx.session_config())?;
         let codec = DistributedCodec::new_combined_with_user(ctx.session_config());
@@ -123,8 +195,13 @@ impl DistributedExec {
                 return exec_err!("Input stage from network boundary was not in Local state");
             };
 
-            let mut spawner =
-                CoordinatorToWorkerTaskSpawner::new(stage, &metrics, &codec, &mut join_set)?;
+            let mut spawner = CoordinatorToWorkerTaskSpawner::new(
+                stage,
+                &metrics,
+                &self.task_metrics,
+                &codec,
+                &mut join_set,
+            )?;
 
             // Right now, we assign random workers to tasks. This might change in the future.
             let start_idx = rand::rng().random_range(0..urls.len());
@@ -135,7 +212,8 @@ impl DistributedExec {
                 workers.push(url.clone());
                 // Spawns the task that feeds this subplan to this worker. There will be as
                 // many as this spawned tasks as workers.
-                let tx = spawner.send_plan_task(Arc::clone(ctx), i, url)?;
+                let (tx, worker_rx) = spawner.send_plan_task(Arc::clone(ctx), i, url)?;
+                spawner.metrics_collection_task(i, worker_rx);
                 spawner.work_unit_feed_task(Arc::clone(ctx), i, tx)?;
             }
 
@@ -185,6 +263,7 @@ impl ExecutionPlan for DistributedExec {
             plan: require_one_child(&children)?,
             prepared_plan: self.prepared_plan.clone(),
             metrics: self.metrics.clone(),
+            task_metrics: Arc::clone(&self.task_metrics),
         }))
     }
 
@@ -254,13 +333,17 @@ struct CoordinatorToWorkerMetrics {
 /// - Building tasks that communicate a serialized plan to multiple workers for further execution.
 /// - Building tasks that stream partition feeds from local [WorkUnitFeedExec] nodes to their
 ///   remote counterparts.
+type WorkerResponseRx =
+    tokio::sync::mpsc::UnboundedReceiver<Result<WorkerToCoordinatorMsg, tonic::Status>>;
+
 struct CoordinatorToWorkerTaskSpawner<'a> {
     plan: &'a Arc<dyn ExecutionPlan>,
     plan_proto: Vec<u8>,
     query_id: Uuid,
     stage_id: usize,
     task_count: usize,
     metrics: &'a CoordinatorToWorkerMetrics,
+    task_metrics: &'a Arc<MetricsStore>,
     join_set: &'a mut JoinSet<Result<()>>,
 }
 
@@ -270,6 +353,7 @@ impl<'a> CoordinatorToWorkerTaskSpawner<'a> {
     fn new(
         stage: &'a LocalStage,
         metrics: &'a CoordinatorToWorkerMetrics,
+        task_metrics: &'a Arc<MetricsStore>,
         codec: &'a dyn PhysicalExtensionCodec,
         join_set: &'a mut JoinSet<Result<()>>,
     ) -> Result<Self> {
@@ -283,18 +367,20 @@ impl<'a> CoordinatorToWorkerTaskSpawner<'a> {
             stage_id: stage.num,
             task_count: stage.tasks,
             metrics,
+            task_metrics,
             join_set,
         })
     }
 
-    /// Instantiates and returns the task sends a serialized plan to specific worker. The returned
-    /// task is just a future that does nothing unless polled.
+    /// Sends a serialized plan to a specific worker and sets up the bidirectional gRPC stream.
+    /// Returns the sender for outbound coordinator-to-worker messages and the receiver for
+    /// inbound worker-to-coordinator messages.
     fn send_plan_task(
         &mut self,
         ctx: Arc<TaskContext>,
         task_i: usize,
         url: Url,
-    ) -> Result<UnboundedSender<CoordinatorToWorkerMsg>> {
+    ) -> Result<(UnboundedSender<CoordinatorToWorkerMsg>, WorkerResponseRx)> {
         let d_cfg = DistributedConfig::from_config_options(ctx.session_config().options())?;
         /// Searches recursively for nodes exposing [crate::WorkUnitFeed]s, and executes their
         /// feeds, keeping into account that some of them might be executed within a
@@ -346,21 +432,25 @@ impl<'a> CoordinatorToWorkerTaskSpawner<'a> {
             &mut work_unit_feed_declarations,
         );
 
+        let task_key = TaskKey {
+            query_id: serialize_uuid(&self.query_id),
+            stage_id: self.stage_id as u64,
+            task_number: task_i as u64,
+        };
         let msg = CoordinatorToWorkerMsg {
             inner: Some(Inner::SetPlanRequest(SetPlanRequest {
                 plan_proto: self.plan_proto.clone(),
                 task_count: self.task_count as u64,
-                task_key: Some(TaskKey {
-                    query_id: serialize_uuid(&self.query_id),
-                    stage_id: self.stage_id as u64,
-                    task_number: task_i as u64,
-                }),
+                task_key: Some(task_key.clone()),
                 work_unit_feed_declarations,
             })),
         };
         let plan_size = self.plan_proto.len();
 
-        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        let (coordinator_to_worker_tx, coordinator_to_worker_rx) =
+            tokio::sync::mpsc::unbounded_channel();
+        let (worker_to_coordinator_tx, worker_to_coordinator_rx) =
+            tokio::sync::mpsc::unbounded_channel();
 
         let channel_resolver = get_distributed_channel_resolver(ctx.as_ref());
 
@@ -370,23 +460,57 @@ impl<'a> CoordinatorToWorkerTaskSpawner<'a> {
         let request = Request::from_parts(
             MetadataMap::from_headers(headers),
             Extensions::default(),
-            futures::stream::once(async { msg }).chain(UnboundedReceiverStream::new(rx)),
+            futures::stream::once(async { msg })
+                .chain(UnboundedReceiverStream::new(coordinator_to_worker_rx)),
         );
 
         let metrics = self.metrics.clone();
+
         self.join_set.spawn(async move {
             let start = Instant::now();
             let mut client = channel_resolver.get_worker_client_for_url(&url).await?;
-            client.coordinator_channel(request).await.map_err(|e| {
+            let response = client.coordinator_channel(request).await.map_err(|e| {
                 tonic_status_to_datafusion_error(&e).unwrap_or_else(|| {
                     exec_datafusion_err!("Error sending plan to worker {url}: {e}")
                 })
             })?;
             metrics.plan_send_latency.record(&start);
             metrics.plan_bytes_sent.add(plan_size);
+            let mut stream = response.into_inner();
+            while let Some(msg) = stream.next().await {
+                if worker_to_coordinator_tx.send(msg).is_err() {
+                    break; // receiver dropped
+                }
+            }
             Ok::<_, DataFusionError>(())
         });
-        Ok(tx)
+
+        Ok((coordinator_to_worker_tx, worker_to_coordinator_rx))
+    }
+
+    /// Receives worker-to-coordinator messages and inserts any collected metrics into the store.
+    /// Runs in a detached spawn so it is not cancelled when the output stream is dropped early.
+    fn metrics_collection_task(
+        &mut self,
+        task_i: usize,
+        mut worker_to_coordinator_rx: WorkerResponseRx,
+    ) {
+        let task_key = TaskKey {
+            query_id: serialize_uuid(&self.query_id),
+            stage_id: self.stage_id as u64,
+            task_number: task_i as u64,
+        };
+        let task_metrics_collection = Arc::clone(self.task_metrics);
+        tokio::spawn(async move {
+            while let Some(Ok(msg)) = worker_to_coordinator_rx.recv().await {
+                let Some(worker_to_coordinator_msg::Inner::TaskMetrics(pre_order_metrics)) =
+                    msg.inner
+                else {
+                    continue;
+                };
+                task_metrics_collection.insert(task_key.clone(), pre_order_metrics.metrics);
+            }
+        });
     }
 
     /// Instantiates and returns the task that based on the different local [WorkUnitFeedExec]

src/execution_plans/mod.rs

@@ -13,7 +13,7 @@ pub mod benchmarks;
 
 pub use broadcast::BroadcastExec;
 pub use children_isolator_union::ChildrenIsolatorUnionExec;
-pub use distributed::DistributedExec;
+pub use distributed::{DistributedExec, MetricsStore};
 pub(crate) use metrics::MetricsWrapperExec;
 pub use network_broadcast::NetworkBroadcastExec;
 pub use network_coalesce::NetworkCoalesceExec;

src/execution_plans/network_broadcast.rs

@@ -2,11 +2,7 @@ use crate::common::require_one_child;
 use crate::distributed_planner::NetworkBoundary;
 use crate::stage::{LocalStage, Stage};
 use crate::worker::WorkerConnectionPool;
-use crate::worker::generated::worker as pb;
-use crate::worker::generated::worker::TaskKey;
-use crate::worker::generated::worker::flight_app_metadata;
 use crate::{BroadcastExec, DistributedTaskContext};
-use dashmap::DashMap;
 use datafusion::common::tree_node::Transformed;
 use datafusion::common::{Result, not_impl_err, plan_err};
 use datafusion::error::DataFusionError;
@@ -125,7 +121,6 @@ pub struct NetworkBroadcastExec {
     pub(crate) properties: Arc<PlanProperties>,
     pub(crate) input_stage: Stage,
     pub(crate) worker_connections: WorkerConnectionPool,
-    pub(crate) metrics_collection: Arc<DashMap<TaskKey, Vec<pb::MetricsSet>>>,
 }
 
 impl NetworkBroadcastExec {
@@ -155,7 +150,6 @@ impl NetworkBroadcastExec {
             properties,
             worker_connections: WorkerConnectionPool::new(0),
             input_stage: Stage::Local(input_stage),
-            metrics_collection: Default::default(),
         }
     }
 
@@ -265,16 +259,7 @@ impl ExecutionPlan for NetworkBroadcastExec {
                 &context,
             )?;
 
-            let metrics_collection = Arc::clone(&self.metrics_collection);
-            let stream = worker_connection.stream_partition(off + partition, move |meta| {
-                if let Some(flight_app_metadata::Content::MetricsCollection(m)) = meta.content {
-                    for task_metrics in m.tasks {
-                        if let Some(task_key) = task_metrics.task_key {
-                            metrics_collection.insert(task_key, task_metrics.metrics);
-                        };
-                    }
-                }
-            })?;
+            let stream = worker_connection.stream_partition(off + partition, |_meta| {})?;
             streams.push(stream);
         }