Finalize worker metrics on coordinator channel EOS

src/coordinator/distributed.rs

@@ -198,7 +198,17 @@ impl ExecutionPlan for DistributedExec {
         let tx = builder.tx();
 
         builder.spawn(async move {
-            let _guard = query_coordinator.end_query_guard();
+            // Dropping this `guard` is what signals the coordinator->worker channel to be dropped,
+            // which triggers a chain reaction that ends up also gracefully closing the
+            // worker->coordinator channel. The flow looks like this:
+            // 1. The query ends normally, as all Arrow RecordBatches are already streamed.
+            // 2. The `guard` here is dropped.
+            // 3. In StageCoordinator::send_plan_task(), `end_stream_notifier` fires and the
+            //    coordinator->worker channel is gracefully ended.
+            // 4. The coordinator->worker channel EOS is received in `impl_coordinator_channel.rs`.
+            // 5. The metrics are send back in the worker->coordinator channel, and then that
+            //    channel is closed.
+            let guard = query_coordinator.end_query_guard();
 
             let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?;
             let result = match d_cfg.dynamic_task_count {
@@ -221,6 +231,7 @@ impl ExecutionPlan for DistributedExec {
                 }
             }
             drop(tx);
+            drop(guard);
             query_coordinator.drain_pending_tasks().await?;
             Ok(())
         });

src/coordinator/query_coordinator.rs

@@ -188,7 +188,14 @@ impl<'a> StageCoordinator<'a> {
                 .map(set_work_unit_send_time)
                 // Keep the request side of the channel open until the query ends: this tail emits
                 // no messages and only completes, once the `Notify` fires. Workers interpret this
-                // EOS of this stream as a query finished/aborted signal.
+                // EOS of this stream as a query finished/aborted signal. The flow looks like this:
+                // 1. The query ends normally, as all Arrow RecordBatches are already streamed.
+                // 2. The end stream notifier guard is dropped in `DistributedExec::execute()`.
+                // 3. Here, `end_stream_notifier` fires and the coordinator->worker channel is
+                //    gracefully ended.
+                // 4. The coordinator->worker channel EOS is received in `impl_coordinator_channel.rs`.
+                // 5. The metrics are send back in the worker->coordinator channel, and then that
+                //    channel is closed.
                 .chain(keep_stream_alive(Arc::clone(self.end_stream_notifier))),
         );
 

src/observability/generated/observability.rs

@@ -15,8 +15,6 @@ pub struct TaskProgress {
     pub task_key: ::core::option::Option<crate::worker::generated::worker::TaskKey>,
     #[prost(uint64, tag = "2")]
     pub total_partitions: u64,
-    #[prost(uint64, tag = "3")]
-    pub completed_partitions: u64,
     #[prost(enumeration = "TaskStatus", tag = "4")]
     pub status: i32,
     #[prost(uint64, tag = "5")]

src/observability/proto/observability.proto

@@ -25,7 +25,7 @@ message TaskKey {
 message TaskProgress {
   TaskKey task_key = 1;
   uint64 total_partitions = 2;
-  uint64 completed_partitions = 3;
+  reserved 3;
   TaskStatus status = 4;
   uint64 output_rows = 5;
 }

src/observability/service.rs

@@ -94,14 +94,11 @@ impl ObservabilityService for ObservabilityServiceImpl {
             // Only include initialized tasks
             if let Some(Ok(task_data)) = task_data_cell.read_now() {
                 let total_partitions = task_data.total_partitions() as u64;
-                let remaining = task_data.num_partitions_remaining() as u64;
-                let completed_partitions = total_partitions.saturating_sub(remaining);
                 let output_rows = output_rows_from_plan(&task_data.base_plan);
 
                 tasks.push(TaskProgress {
                     task_key: Some((*internal_key).clone()),
                     total_partitions,
-                    completed_partitions,
                     status: TaskStatus::Running as i32,
                     output_rows,
                 });

src/worker/impl_coordinator_channel.rs

@@ -1,28 +1,32 @@
-use crate::common::deserialize_uuid;
+use crate::common::{TreeNodeExt, deserialize_uuid};
 use crate::execution_plans::SamplerExec;
+use crate::metrics::proto::df_metrics_set_to_proto;
+use crate::protobuf::datafusion_error_to_tonic_status;
 use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time};
 use crate::worker::LocalWorkerContext;
 use crate::worker::generated::worker::coordinator_to_worker_msg::Inner;
 use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration;
 use crate::worker::generated::worker::worker_service_server::WorkerService;
 use crate::worker::generated::worker::{
-    CoordinatorToWorkerMsg, WorkerToCoordinatorMsg, worker_to_coordinator_msg,
+    CoordinatorToWorkerMsg, TaskMetrics, WorkerToCoordinatorMsg, worker_to_coordinator_msg,
 };
 use crate::worker::task_data::TaskDataMetrics;
 use crate::{
     DistributedCodec, DistributedConfig, DistributedExt, DistributedTaskContext, TaskData, Worker,
     WorkerQueryContext,
 };
 use datafusion::common::DataFusionError;
+use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::execution::SessionStateBuilder;
+use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionConfig;
 use datafusion_proto::physical_plan::AsExecutionPlan;
 use datafusion_proto::protobuf::PhysicalPlanNode;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt, TryStreamExt};
-use std::sync::atomic::AtomicUsize;
 use std::sync::{Arc, OnceLock};
 use tokio::sync::oneshot;
+use tokio::sync::oneshot::Sender;
 use tonic::{Request, Response, Status, Streaming};
 use url::Url;
 
@@ -105,12 +109,10 @@ impl Worker {
                 SamplerExec::kick_off_first_sampler(Arc::clone(&plan), Arc::clone(&task_ctx))?;
 
             // Initialize partition count to the number of partitions in the stage
-            let total_partitions = plan.properties().partitioning.partition_count();
             Ok::<_, DataFusionError>(TaskData {
                 base_plan: plan,
                 final_plan: Arc::new(OnceLock::new()),
                 task_ctx,
-                num_partitions_remaining: Arc::new(AtomicUsize::new(total_partitions)),
                 metrics_tx: match collect_metrics {
                     true => Arc::new(std::sync::Mutex::new(Some(metrics_tx))),
                     false => Arc::new(std::sync::Mutex::new(None)),
@@ -119,15 +121,33 @@ impl Worker {
             })
         };
 
-        entry.write(task_data().await.map_err(Arc::new)).map_err(|_| {
-            Status::internal(format!(
+        let task_data_result = task_data().await.map_err(Arc::new);
+
+        entry
+            .write(task_data_result.clone())
+            .map_err(|_| Status::internal(format!(
                 "Logic error while setting plan for TaskKey {key:?}: the plan was set twice. This is a bug in datafusion-distributed, please report it."
-            ))
-        })?;
+            )))?;
+
+        let task_data = task_data_result
+            .map_err(DataFusionError::Shared)
+            .map_err(datafusion_error_to_tonic_status)?;
 
         // Continue reading remaining messages (work unit feed data) in the background.
         let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders);
         let task_data_entries = Arc::clone(&self.task_data_entries);
+
+        // This tokio task takes ownership of the `oneshot::Sender<pb::TaskMetrics>` that keeps
+        // alive the worker->coordinator stream. as soon as this task ends, the runtime metrics
+        // are send back and the worker->coordinator stream ends. The flow is the following:
+        // 1. The query ends normally, as all Arrow RecordBatches are already streamed.
+        // 2. In DistributedExec::execute(), the end query guard is dropped.
+        // 3. In StageCoordinator::send_plan_task(), `end_stream_notifier` fires and the
+        //    coordinator->worker channel is gracefully ended.
+        // 4. The coordinator->worker channel EOS is received by this same function, ending the
+        //    while loop inside this `tokio::spawn` below.
+        // 5. The metrics are send back in the worker->coordinator channel, and then that channel
+        //    is closed.
         #[allow(clippy::disallowed_methods)]
         tokio::spawn(async move {
             let mut body = body.map_ok(set_work_unit_received_time);
@@ -173,8 +193,20 @@ impl Worker {
                     }
                 }
             }
-            #[allow(clippy::disallowed_methods)]
-            tokio::spawn(async move { task_data_entries.invalidate(&key).await });
+
+            let metrics_tx = task_data.metrics_tx.lock().unwrap().take();
+            if let Some(Ok(plan)) = task_data.final_plan.get() {
+                let d_ctx = DistributedTaskContext {
+                    task_index: key.task_number as usize,
+                    task_count: request.task_count as usize,
+                };
+                let task_data_metrics = &task_data.task_data_metrics;
+                task_data_metrics.mark_execution_finished();
+                if let Some(metrics_tx) = metrics_tx {
+                    send_metrics_via_channel(metrics_tx, plan, d_ctx, task_data_metrics);
+                }
+            }
+            task_data_entries.invalidate(&key).await
         });
 
         let load_info_stream = FuturesUnordered::from_iter(load_info_rxs)
@@ -191,9 +223,9 @@ impl Worker {
                 })
             }));
 
-        // Stream back the metrics once the task finishes executing.
-        // The oneshot receiver resolves when impl_execute_task sends the collected
-        // metrics after all partitions have finished or been dropped.
+        // Stream back metrics when the coordinator channel reaches EOS. At that point the
+        // coordinator has closed the query-scoped request stream, so any remaining task state can
+        // be finalized even if some partition streams were not dropped through the normal path.
         let metrics_stream = metrics_rx.into_stream();
         let metrics_stream = metrics_stream.filter_map(async |task_metrics_or_channel_dropped| {
             let task_metrics = task_metrics_or_channel_dropped.ok()?;
@@ -211,3 +243,28 @@ impl Worker {
 fn missing(field: &'static str) -> impl FnOnce() -> Status {
     move || Status::invalid_argument(format!("Missing field '{field}'"))
 }
+
+/// Collects metrics from the plan in pre-order traversal order and sends them via the
+/// coordinator channel oneshot.
+fn send_metrics_via_channel(
+    metrics_tx: Sender<TaskMetrics>,
+    plan: &Arc<dyn ExecutionPlan>,
+    dt_ctx: DistributedTaskContext,
+    task_data_metrics: &Arc<TaskDataMetrics>,
+) {
+    let mut pre_order_plan_metrics = vec![];
+    let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| {
+        pre_order_plan_metrics.push(
+            node.metrics()
+                .and_then(|m| df_metrics_set_to_proto(&m).ok())
+                .unwrap_or_default(),
+        );
+        Ok(TreeNodeRecursion::Continue)
+    });
+
+    // Ignore send errors — the coordinator channel may have been dropped (e.g. query cancelled).
+    let _ = metrics_tx.send(TaskMetrics {
+        pre_order_plan_metrics,
+        task_metrics: Some(task_data_metrics.to_proto_metrics_set()),
+    });
+}

src/worker/impl_execute_task.rs

@@ -1,34 +1,26 @@
-use crate::common::{TreeNodeExt, now_ns, on_drop_stream};
-use crate::metrics::proto::df_metrics_set_to_proto;
+use crate::DistributedConfig;
+use crate::common::now_ns;
 use crate::protobuf::datafusion_error_to_tonic_status;
-use crate::worker::generated::worker::{FlightAppMetadata, TaskMetrics};
+use crate::worker::generated::worker::ExecuteTaskRequest;
+use crate::worker::generated::worker::FlightAppMetadata;
+use crate::worker::generated::worker::worker_service_server::WorkerService;
+use crate::worker::spawn_select_all::spawn_select_all;
 use crate::worker::worker_service::{TaskDataEntries, Worker};
-use crate::{DistributedConfig, DistributedTaskContext};
 use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder};
 use arrow_flight::error::FlightError;
 use arrow_select::dictionary::garbage_collect_any_dictionary;
 use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions};
-use datafusion::common::tree_node::TreeNodeRecursion;
-use datafusion::common::{Result, exec_err, internal_err};
-
-use crate::worker::generated::worker::ExecuteTaskRequest;
-use crate::worker::generated::worker::worker_service_server::WorkerService;
-use crate::worker::spawn_select_all::spawn_select_all;
-use crate::worker::task_data::TaskDataMetrics;
 use datafusion::arrow::ipc::CompressionType;
 use datafusion::arrow::ipc::writer::IpcWriteOptions;
 use datafusion::common::exec_datafusion_err;
+use datafusion::common::{Result, exec_err, internal_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
-use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use futures::TryStreamExt;
 use prost::Message;
 use std::sync::Arc;
-use std::sync::Mutex;
-use std::sync::atomic::Ordering;
 use std::time::Duration;
-use tokio::sync::oneshot::Sender;
 use tokio_stream::StreamExt;
 use tonic::{Request, Response, Status};
 
@@ -66,10 +58,7 @@ pub(crate) async fn execute_local_task(
 
     let plan = task_data.plan(producer_head)?;
     let task_ctx = task_data.task_ctx;
-    let d_cfg = DistributedConfig::from_config_options(task_ctx.session_config().options())?;
-    let d_ctx = *DistributedTaskContext::from_ctx(&task_ctx).as_ref();
 
-    let send_metrics = d_cfg.collect_metrics;
     let partition_count = plan.properties().partitioning.partition_count();
     let plan_name = plan.name();
 
@@ -87,29 +76,6 @@ pub(crate) async fn execute_local_task(
         let stream = plan.execute(partition as usize, Arc::clone(&task_ctx))?;
         let stream_schema = plan.schema();
 
-        let plan = Arc::clone(&plan);
-
-        let task_data_entries = Arc::clone(task_data_entries);
-        let num_partitions_remaining = Arc::clone(&task_data.num_partitions_remaining);
-        let metrics_tx = Arc::clone(&task_data.metrics_tx);
-        let task_data_metrics = Arc::clone(&task_data.task_data_metrics);
-        let key = key.clone();
-        let stream = on_drop_stream(stream, move || {
-            // Stream was dropped before fully consumed -- see https://github.com/datafusion-contrib/datafusion-distributed/issues/412
-            // Send metrics via the coordinator channel so they are not lost.
-            if num_partitions_remaining.fetch_sub(1, Ordering::SeqCst) == 1 {
-                // Fire-and-forget background tokio task to handle async
-                // invalidate() within synchronous on_drop_stream.
-                #[allow(clippy::disallowed_methods)]
-                tokio::spawn(async move {
-                    task_data_entries.invalidate(&key).await;
-                });
-                task_data_metrics.mark_execution_finished();
-                if send_metrics {
-                    send_metrics_via_channel(&metrics_tx, &plan, d_ctx, &task_data_metrics);
-                }
-            }
-        });
         streams.push(Box::pin(RecordBatchStreamAdapter::new(stream_schema, stream)) as _);
     }
     Ok((streams, task_ctx))
@@ -208,39 +174,6 @@ fn build_flight_data_stream(
     Ok(stream)
 }
 
-/// Collects metrics from the plan in pre-order traversal order and sends them via the
-/// coordinator channel oneshot.
-fn send_metrics_via_channel(
-    metrics_tx: &Arc<Mutex<Option<Sender<TaskMetrics>>>>,
-    plan: &Arc<dyn ExecutionPlan>,
-    dt_ctx: DistributedTaskContext,
-    task_data_metrics: &Arc<TaskDataMetrics>,
-) {
-    let mut pre_order_plan_metrics = vec![];
-    let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| {
-        pre_order_plan_metrics.push(
-            node.metrics()
-                .and_then(|m| df_metrics_set_to_proto(&m).ok())
-                .unwrap_or_default(),
-        );
-        Ok(TreeNodeRecursion::Continue)
-    });
-
-    let tx = {
-        let mut guard = match metrics_tx.lock() {
-            Ok(g) => g,
-            Err(_) => return,
-        };
-        guard.take()
-    };
-    let Some(tx) = tx else { return };
-    // Ignore send errors — the coordinator channel may have been dropped (e.g. query cancelled).
-    let _ = tx.send(TaskMetrics {
-        pre_order_plan_metrics,
-        task_metrics: Some(task_data_metrics.to_proto_metrics_set()),
-    });
-}
-
 /// Garbage collects values sub-arrays.
 ///
 /// We apply this before sending RecordBatches over the network to avoid sending

src/worker/task_data.rs

@@ -8,7 +8,6 @@ use datafusion::execution::TaskContext;
 use datafusion::physical_expr_common::metrics::CustomMetricValue;
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use std::sync::Arc;
-use std::sync::atomic::{AtomicUsize, Ordering};
 use std::time::Duration;
 use tokio::sync::oneshot;
 
@@ -20,12 +19,6 @@ pub struct TaskData {
     pub(super) task_ctx: Arc<TaskContext>,
     pub(crate) base_plan: Arc<dyn ExecutionPlan>,
     pub(crate) final_plan: Arc<OnceLockResult<Arc<dyn ExecutionPlan>>>,
-    /// `num_partitions_remaining` is initialized to the total number of partitions in the task (not
-    /// only tasks in the partition group). This is decremented for each request to the endpoint
-    /// for this task. Once this count is zero, the task is likely complete. The task may not be
-    /// complete because it's possible that the same partition was retried and this count was
-    /// decremented more than once for the same partition.
-    pub(super) num_partitions_remaining: Arc<AtomicUsize>,
     /// Sender half of the metrics channel. `impl_execute_task` takes this (via `Option::take`)
     /// once all partitions have finished or been dropped, sending the collected metrics back to
     /// the coordinator through the `CoordinatorChannel` side channel.
@@ -110,11 +103,6 @@ impl TaskDataMetrics {
 }
 
 impl TaskData {
-    /// Returns the number of partitions remaining to be processed.
-    pub(crate) fn num_partitions_remaining(&self) -> usize {
-        self.num_partitions_remaining.load(Ordering::SeqCst)
-    }
-
     /// Returns the total number of partitions in this task.
     pub(crate) fn total_partitions(&self) -> usize {
         match self.final_plan.get() {
@@ -135,13 +123,7 @@ impl TaskData {
             let producer_head =
                 ProducerHead::from_proto(producer_head, &self.base_plan.schema(), &self.task_ctx)?;
 
-            let plan = producer_head.insert(Arc::clone(&self.base_plan))?;
-
-            self.num_partitions_remaining.store(
-                plan.output_partitioning().partition_count(),
-                Ordering::SeqCst,
-            );
-            Ok(plan)
+            Ok(producer_head.insert(Arc::clone(&self.base_plan))?)
         });
         match result {
             Ok(plan) => Ok(Arc::clone(plan)),