From 6166ac9877521c73392c4a25b74adb5aee20b365 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 2 Jul 2026 09:22:08 +0200 Subject: [PATCH 1/2] Fix potential deadlock in worker->coordinator channels --- src/coordinator/distributed.rs | 13 +++- src/coordinator/query_coordinator.rs | 9 ++- src/worker/impl_coordinator_channel.rs | 82 ++++++++++++++++++++++---- src/worker/impl_execute_task.rs | 81 +++---------------------- 4 files changed, 98 insertions(+), 87 deletions(-) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index d7d62a08..2c4bc611 100644 --- a/src/coordinator/distributed.rs +++ b/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(()) }); diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 6f0010e6..2e0f2bf0 100644 --- a/src/coordinator/query_coordinator.rs +++ b/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))), ); diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index fcb3944a..b937de73 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,12 +1,14 @@ -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::{ @@ -14,7 +16,9 @@ use crate::{ 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; @@ -23,6 +27,7 @@ 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; @@ -119,15 +124,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` 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 +196,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 +226,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 +246,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, + plan: &Arc, + dt_ctx: DistributedTaskContext, + task_data_metrics: &Arc, +) { + 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()), + }); +} diff --git a/src/worker/impl_execute_task.rs b/src/worker/impl_execute_task.rs index 9be4e51b..ef3b9e2f 100644 --- a/src/worker/impl_execute_task.rs +++ b/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>>>, - plan: &Arc, - dt_ctx: DistributedTaskContext, - task_data_metrics: &Arc, -) { - 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 From 2b108ea82319d2ad7599057f6ec358a394bbc880 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Thu, 2 Jul 2026 19:51:57 +0200 Subject: [PATCH 2/2] Remove num partitions --- src/observability/generated/observability.rs | 2 -- src/observability/proto/observability.proto | 2 +- src/observability/service.rs | 3 --- src/worker/impl_coordinator_channel.rs | 3 --- src/worker/task_data.rs | 20 +------------------- src/worker/test_utils/worker_handles.rs | 6 ------ 6 files changed, 2 insertions(+), 34 deletions(-) diff --git a/src/observability/generated/observability.rs b/src/observability/generated/observability.rs index c8da3372..ee40102b 100644 --- a/src/observability/generated/observability.rs +++ b/src/observability/generated/observability.rs @@ -15,8 +15,6 @@ pub struct TaskProgress { pub task_key: ::core::option::Option, #[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")] diff --git a/src/observability/proto/observability.proto b/src/observability/proto/observability.proto index 07925d54..f403aa85 100644 --- a/src/observability/proto/observability.proto +++ b/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; } diff --git a/src/observability/service.rs b/src/observability/service.rs index 7b7c5473..f49c4270 100644 --- a/src/observability/service.rs +++ b/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, }); diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index b937de73..c3beb3e0 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -24,7 +24,6 @@ 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; @@ -110,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)), diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index 97b2e806..c29d60aa 100644 --- a/src/worker/task_data.rs +++ b/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, pub(crate) base_plan: Arc, pub(crate) final_plan: Arc>>, - /// `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, /// 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)), diff --git a/src/worker/test_utils/worker_handles.rs b/src/worker/test_utils/worker_handles.rs index 477b1de5..f328145f 100644 --- a/src/worker/test_utils/worker_handles.rs +++ b/src/worker/test_utils/worker_handles.rs @@ -11,7 +11,6 @@ use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ExecutionPlan; use hyper_util::rt::TokioIo; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; use tokio::net::TcpListener; use tonic::transport::{Endpoint, Server}; use url::Url; @@ -101,13 +100,11 @@ impl MemoryWorkerHandle { None, )?; let plan = build_plan(input)?; - let partition_count = plan.properties().partitioning.partition_count(); register_plan_on_worker( &self.worker, task_ctx, plan, test_task_key_with_query(query_id, self.task_index as _), - partition_count, ) .await; Ok(()) @@ -173,7 +170,6 @@ impl TcpWorkerHandle { task_ctx, plan, test_task_key_with_query(query_id, self.task_index as _), - self.partitions_batches.len(), ) .await; Ok(()) @@ -205,7 +201,6 @@ pub async fn register_plan_on_worker( task_ctx: Arc, plan: Arc, task_key: TaskKey, - partition_count: usize, ) { let swmr_task_data = worker .task_data_entries @@ -217,7 +212,6 @@ pub async fn register_plan_on_worker( task_ctx, base_plan: plan, final_plan: Default::default(), - num_partitions_remaining: Arc::new(AtomicUsize::new(partition_count)), metrics_tx: Arc::new(std::sync::Mutex::new(Some(metrics_tx))), task_data_metrics: Arc::new(TaskDataMetrics::new(0)), }))