From 7316d6155c28e3e9cd9852af3fdfa52850bcb220 Mon Sep 17 00:00:00 2001 From: Sergio Esteves Date: Wed, 23 Sep 2026 14:09:59 +0100 Subject: [PATCH 1/2] fix: finalize metrics for unexecuted distributed tasks --- Cargo.lock | 1 + Cargo.toml | 1 + docs/source/user-guide/05-metrics.md | 58 ++++- src/coordinator/distributed.rs | 217 +++++++++++++++++- src/coordinator/mod.rs | 2 +- src/coordinator/query_coordinator.rs | 146 +++++++++++- src/coordinator/store.rs | 134 ++++++++++- src/distributed_planner/distributed_config.rs | 6 + .../distributed_query_planner.rs | 6 +- src/distributed_planner/mod.rs | 1 + src/execution_plans/sampler.rs | 39 ++++ src/lib.rs | 2 +- src/metrics/mod.rs | 4 + src/metrics/task_metrics_rewriter.rs | 98 +++++++- src/protocol/worker_channel.rs | 10 + src/worker/impl_coordinator_channel.rs | 161 +++++++++++-- src/worker/task_data.rs | 23 +- tests/metrics_collection.rs | 116 ++++++++++ 18 files changed, 962 insertions(+), 63 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8fa95485e..08350241b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2389,6 +2389,7 @@ dependencies = [ "hyper-util", "insta", "itertools 0.14.0", + "log", "moka", "num-traits", "parquet", diff --git a/Cargo.toml b/Cargo.toml index f98cd3edd..26d426a49 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ tokio = { version = "1.48", features = ["full"] } http = "1.3.1" itertools = "0.14.0" futures = "0.3.31" +log = "0.4.27" url = "2.5.7" uuid = "1.19" delegate = "0.13.4" diff --git a/docs/source/user-guide/05-metrics.md b/docs/source/user-guide/05-metrics.md index 1bdb2defd..062b06cad 100644 --- a/docs/source/user-guide/05-metrics.md +++ b/docs/source/user-guide/05-metrics.md @@ -8,21 +8,35 @@ display them. Distributed DataFusion does this for you, and exposes two functions you can use to build your own EXPLAIN ANALYZE in application code. -## Enabling collection +## Configuring collection -Metrics collection across network boundaries is **on by default**. You can toggle it explicitly: +Metrics collection across network boundaries is **on by default**. To disable it, call +`.with_distributed_metrics_collection(false)?` on the session-state builder. + +When enabled, each worker streams the metrics for its tasks back to the coordinator on a dedicated +channel, so they are not lost even if the result stream is dropped early (for example by a `LIMIT`). + +Post-query cleanup, metrics collection, and completed dynamic-filter collection share one deadline +starting when the result stream ends. The default is five seconds. To adjust it, configure +`DistributedConfig::metrics_finalization_timeout_ms` before planning the query: ```rust +use datafusion::execution::SessionStateBuilder; +use datafusion_distributed::{DistributedConfig, DistributedExt, SessionStateBuilderExt}; + let state = SessionStateBuilder::new() .with_default_features() + .with_distributed_option_extension(DistributedConfig { + metrics_finalization_timeout_ms: 10_000, + ..Default::default() + }) .with_distributed_worker_resolver(/* ... */) .with_distributed_planner() - .with_distributed_metrics_collection(true) // default is true .build(); ``` -When enabled, each worker streams the metrics for its tasks back to the coordinator on a dedicated -channel, so they are not lost even if the result stream is dropped early (for example by a `LIMIT`). +Set it to zero to avoid waiting for late reports. An incomplete metrics result still returns an +error rather than treating missing work as zero. ## Rendering a plan with metrics @@ -33,8 +47,9 @@ These functions, all exported from the crate root, do the work: used to execute the plan. When displaying both dynamic filters and metrics, apply the dynamic-filter rewrite first. - `rewrite_distributed_plan_with_metrics(plan, format)` — folds every task's metrics back into the - coordinator's copy of the plan. It waits for all worker metrics to arrive, so the result is always - complete. The `format` is a `DistributedMetricsFormat`: + coordinator's copy of the plan. It waits for complete worker reports and returns an error if a + channel closes without a report or finalization exceeds the configured deadline. The `format` is a + `DistributedMetricsFormat`: - `Aggregated` — metrics from all tasks of a stage are summed/aggregated into one value per node. - `PerTask` — each metric collects its per-task values into a map keyed by task id (`output_rows={0:.., 1:..}`) so you can see each task individually. @@ -92,3 +107,32 @@ runtime metrics, including network-level metrics on the boundaries: > If `plan` is not a distributed plan (its root is not a `DistributedExec`), > `rewrite_distributed_plan_with_metrics` returns it unchanged, so it is always safe to call. + +## Inspecting a partial snapshot + +`DistributedExec::metrics_snapshot()` returns the metrics received **so far** without waiting. +It returns `None` if collection is disabled or the plan has not yet been prepared. The snapshot +has separate `pending`, `missing_reports`, and `reported` task keys; `all_terminal()` says whether +every expected task has reported or its stream has ended. `all_reported()` says whether metrics +are available for every task; use it when accounting must be complete. `unexecuted_tasks()` identifies planned +tasks that reported actual AQE sampling metrics but never received an `ExecuteTask` call. These +tasks have a `task_not_executed=1` stage metric instead of fabricated execution timestamps. + +```rust +use datafusion_distributed::DistributedExec; + +if let Some(exec) = plan.downcast_ref::() { + if let Some(snapshot) = exec.metrics_snapshot() { + println!("{} reports; {} pending", snapshot.reported.len(), snapshot.pending.len()); + if !snapshot.all_reported() { + // Do not treat a missing task report as zero I/O. + println!("missing reports: {:?}", snapshot.missing_reports); + } + } +} +``` + +The existing `DistributedExec::wait_for_metrics()` only returns `Some` for a complete set of +reports. If any task fails to report or the wait times out, it returns `None`; call +`metrics_snapshot()` to inspect what arrived. A closed stream without a metrics report is not +equivalent to an unexecuted task with a report. diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index e2b2508b7..ecb34f26b 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -2,14 +2,15 @@ use crate::common::require_one_child; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; -use crate::coordinator::store::{Store, task_keys_for_plan}; +use crate::coordinator::store::{Store, StoreSnapshot, task_keys_for_plan}; +use crate::distributed_planner::DEFAULT_METRICS_FINALIZATION_TIMEOUT_MS; use crate::dynamic_filtering::{ is_dynamic_filtering_enabled, sever_dynamic_filter_relationships_in_plan_for_display, }; use crate::{DistributedConfig, TaskCompletedDynamicFilters, TaskKey, TaskMetrics}; use datafusion::common::internal_datafusion_err; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{HashMap, Result, exec_err}; +use datafusion::common::{HashMap, Result, exec_datafusion_err, exec_err}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_common::metrics::MetricsSet; @@ -19,6 +20,49 @@ use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, Pla use futures::StreamExt; use std::fmt::Formatter; use std::sync::{Arc, OnceLock}; +use std::time::Duration; +use tokio::sync::watch; +use tokio::time::Instant; + +/// Non-blocking, point-in-time view of the reports expected from a prepared distributed plan. +#[derive(Debug)] +pub struct TaskMetricsSnapshot { + /// Metrics received so far, including actual sampling metrics from tasks that never ran. + pub reported: HashMap, + /// Tasks whose coordinator channels are still open and have not reported metrics. + pub pending: Vec, + /// Tasks whose channels closed without a report. Their work cannot be accounted for. + pub missing_reports: Vec, +} + +impl TaskMetricsSnapshot { + fn from_store(snapshot: StoreSnapshot) -> Self { + Self { + reported: snapshot.reported, + pending: snapshot.pending, + missing_reports: snapshot.terminal_without_report, + } + } + + /// Whether every expected task has either reported or had its stream end without a report. + pub fn all_terminal(&self) -> bool { + self.pending.is_empty() + } + + /// Whether every expected task has actually reported metrics. This distinguishes a + /// finished-but-incomplete query from one whose measurements are complete. + pub fn all_reported(&self) -> bool { + self.all_terminal() && self.missing_reports.is_empty() + } + + /// Tasks that reported real sampling metrics but never received ExecuteTask. + pub fn unexecuted_tasks(&self) -> Vec { + self.reported + .iter() + .filter_map(|(key, report)| report.was_not_executed().then_some(*key)) + .collect() + } +} /// [ExecutionPlan] that executes the inner plan in distributed mode. /// Before executing it, two modifications are lazily performed on the plan: @@ -50,6 +94,9 @@ pub struct DistributedExec { pub(crate) metrics_store: Option>>, /// Storage for the completed dynamic filters reported by each worker task. pub(crate) completed_dynamic_filter_store: Option>>, + /// Set when the result stream ends; all post-query waits use the same deadline. + finalization_deadline: watch::Sender>, + finalization_timeout: Duration, } #[derive(Debug, Clone)] @@ -62,12 +109,15 @@ pub(super) struct PreparedPlan { impl DistributedExec { pub fn new(base_plan: Arc) -> Self { + let (finalization_deadline, _) = watch::channel(None); Self { base_plan, prepared_plan: Arc::new(OnceLock::new()), metrics: ExecutionPlanMetricsSet::new(), metrics_store: None, completed_dynamic_filter_store: None, + finalization_deadline, + finalization_timeout: Duration::from_millis(DEFAULT_METRICS_FINALIZATION_TIMEOUT_MS), } } @@ -89,12 +139,41 @@ impl DistributedExec { self } - /// Waits until all worker tasks have reported their metrics back via the coordinator channel - /// if metrics collection is enabled. + pub(crate) fn with_finalization_timeout(mut self, timeout: Duration) -> Self { + self.finalization_timeout = timeout; + self + } + + async fn finalization_deadline(&self) -> Instant { + wait_for_finalization_deadline(self.finalization_deadline.subscribe()).await + } + + /// Waits for complete metrics, if collection is enabled and execution has been prepared. + /// Returns `None` if a task failed to report or the finalization timeout elapsed; use + /// [`Self::metrics_snapshot`] to inspect reports and their completeness in that case. pub async fn wait_for_metrics(&self) -> Option> { - let task_metrics = self.metrics_store.as_ref()?; + self.complete_metrics().await.ok() + } + + /// Returns all available metrics immediately, including gaps in task numbers, along with + /// pending and terminal-without-report task keys. Returns `None` if metrics are disabled or + /// the distributed plan has not yet been prepared. + pub fn metrics_snapshot(&self) -> Option { + let store = self.metrics_store.as_ref()?; let plan = &self.prepared_plan.get()?.plan_for_viz; - Some(task_metrics.wait_for(&task_keys_for_plan(plan)).await) + Some(TaskMetricsSnapshot::from_store( + store.snapshot(&task_keys_for_plan(plan)), + )) + } + + pub(crate) async fn complete_metrics(&self) -> Result> { + let store = self + .metrics_store + .as_ref() + .ok_or_else(|| exec_datafusion_err!("metrics collection is disabled"))?; + let plan = &self.prepared_plan()?.plan_for_viz; + let keys = task_keys_for_plan(plan); + wait_for_complete_metrics(store, &keys, self.finalization_deadline().await).await } /// Waits until all worker tasks have reported their completed dynamic filters back via @@ -104,7 +183,12 @@ impl DistributedExec { ) -> Option> { let store = self.completed_dynamic_filter_store.as_ref()?; let plan = &self.prepared_plan.get()?.plan_for_viz; - Some(store.wait_for(&task_keys_for_plan(plan)).await) + tokio::time::timeout_at( + self.finalization_deadline().await, + store.wait_for(&task_keys_for_plan(plan)), + ) + .await + .ok() } fn prepared_plan(&self) -> Result { @@ -152,10 +236,41 @@ impl DistributedExec { metrics: self.metrics.clone(), metrics_store: self.metrics_store.clone(), completed_dynamic_filter_store: self.completed_dynamic_filter_store.clone(), + finalization_deadline: self.finalization_deadline.clone(), + finalization_timeout: self.finalization_timeout, })) } } +async fn wait_for_finalization_deadline(mut rx: watch::Receiver>) -> Instant { + (*rx.wait_for(Option::is_some) + .await + .expect("DistributedExec owns the finalization deadline sender")) + .expect("deadline was set before the receiver woke") +} + +async fn wait_for_complete_metrics( + store: &Store, + keys: &[TaskKey], + deadline: Instant, +) -> Result> { + let snapshot = tokio::time::timeout_at(deadline, store.wait_for_terminal(keys)) + .await + .map_err(|_| { + exec_datafusion_err!( + "timed out waiting for worker task metrics; pending: {:?}", + store.snapshot(keys).pending + ) + })?; + if !snapshot.terminal_without_report.is_empty() { + return exec_err!( + "worker task metrics missing after coordinator channel closed: {:?}", + snapshot.terminal_without_report + ); + } + Ok(snapshot.reported) +} + impl DisplayAs for DistributedExec { fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { write!(f, "DistributedExec") @@ -200,6 +315,8 @@ impl ExecutionPlan for DistributedExec { metrics: self.metrics.clone(), metrics_store: self.metrics_store.clone(), completed_dynamic_filter_store: self.completed_dynamic_filter_store.clone(), + finalization_deadline: self.finalization_deadline.clone(), + finalization_timeout: self.finalization_timeout, })) } @@ -221,6 +338,8 @@ impl ExecutionPlan for DistributedExec { let base_plan = Arc::clone(&self.base_plan); let prepared_plan = Arc::clone(&self.prepared_plan); let collect_dynamic_filters = self.completed_dynamic_filter_store.is_some(); + let finalization_deadline = self.finalization_deadline.clone(); + let finalization_timeout = self.finalization_timeout; let query_coordinator = Arc::new(QueryCoordinator::new( Arc::clone(&context), @@ -243,7 +362,8 @@ impl ExecutionPlan for DistributedExec { // 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 guard = query_coordinator + .end_query_guard(finalization_deadline.clone(), finalization_timeout); let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?; let mut prepared = match d_cfg.dynamic_task_count { @@ -271,7 +391,16 @@ impl ExecutionPlan for DistributedExec { } drop(guard); drop(tx); - query_coordinator.drain_pending_tasks().await?; + // DataFusion does not close the result stream until this spawned task returns. + // A half-open coordinator channel can otherwise hang result collection itself, + // even if the later metrics rewrite has its own timeout. Cancel stalled + // background tasks on timeout; their metric receivers then become terminal + // without a report, so accounting remains explicitly incomplete. + query_coordinator + .drain_pending_tasks( + wait_for_finalization_deadline(finalization_deadline.subscribe()).await, + ) + .await?; Ok(()) }); @@ -282,3 +411,73 @@ impl ExecutionPlan for DistributedExec { Some(self.metrics.clone_inner()) } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::physical_plan::metrics::MetricsSet; + use uuid::Uuid; + + #[tokio::test] + async fn a_never_closing_channel_cannot_block_complete_metrics_indefinitely() { + let store = Store::::new(); + let key = TaskKey { + query_id: Uuid::new_v4(), + stage_id: 1, + task_number: 0, + }; + let deadline = Instant::now() + Duration::from_millis(25); + let result = wait_for_complete_metrics(&store, &[key], deadline).await; + assert!( + result + .unwrap_err() + .to_string() + .contains("timed out waiting") + ); + assert_eq!(store.snapshot(&[key]).pending, vec![key]); + // Rewriting after the drain used up the budget must not start another wait. + let result = tokio::time::timeout( + Duration::from_millis(100), + wait_for_complete_metrics(&store, &[key], deadline), + ) + .await + .expect("expired deadline should return without another timeout window"); + assert!(result.is_err()); + } + + #[tokio::test] + async fn closed_channel_without_metrics_is_an_error_not_zero_metrics() { + let store = Store::::new(); + let query_id = Uuid::new_v4(); + let reported = TaskKey { + query_id, + stage_id: 1, + task_number: 0, + }; + let lost = TaskKey { + query_id, + stage_id: 1, + task_number: 2, + }; + store.insert( + reported, + TaskMetrics { + pre_order_plan_metrics: vec![], + task_metrics: MetricsSet::new(), + }, + ); + store.mark_terminal(lost); + let result = wait_for_complete_metrics( + &store, + &[reported, lost], + Instant::now() + Duration::from_secs(1), + ) + .await; + assert!(result.unwrap_err().to_string().contains("metrics missing")); + let snapshot = TaskMetricsSnapshot::from_store(store.snapshot(&[reported, lost])); + assert!(snapshot.all_terminal()); + assert!(!snapshot.all_reported()); + assert_eq!(snapshot.missing_reports, vec![lost]); + assert_eq!(snapshot.reported.len(), 1); + } +} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index 746052a6f..b41321ce3 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -6,6 +6,6 @@ mod prepare_static_plan; mod query_coordinator; mod store; -pub use distributed::DistributedExec; +pub use distributed::{DistributedExec, TaskMetricsSnapshot}; pub(crate) use dynamic_filter_registry::DynamicFilterRegistry; pub(crate) use store::Store; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 4889f888f..1a8c24c1b 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -36,8 +36,10 @@ use datafusion::prelude::SessionConfig; use futures::{Stream, StreamExt, TryStreamExt}; use std::ops::DerefMut; use std::sync::{Arc, Mutex}; -use tokio::sync::Notify; +use std::time::Duration; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; +use tokio::sync::{Notify, watch}; +use tokio::time::Instant as TokioInstant; use tokio_stream::wrappers::UnboundedReceiverStream; use url::Url; use uuid::Uuid; @@ -109,16 +111,33 @@ impl QueryCoordinator { /// returns a guard that, when dropped, it signals all the coordinator->worker connections that /// the query is finished, ending them, and propagating the EOS to the workers so that they can /// clean up any remaining state. - pub(super) fn end_query_guard(&self) -> NotifyGuard { - NotifyGuard(Arc::clone(&self.end_stream_notifier)) + pub(super) fn end_query_guard( + &self, + deadline: watch::Sender>, + timeout: Duration, + ) -> NotifyGuard { + NotifyGuard { + notifier: Arc::clone(&self.end_stream_notifier), + deadline, + timeout, + } } /// Blocks until all background tasks have finished (e.g., sending WorkUnit feeds, or collecting /// metrics) - pub(super) async fn drain_pending_tasks(self: Arc) -> Result<()> { + pub(super) async fn drain_pending_tasks(self: Arc, deadline: TokioInstant) -> Result<()> { let join_set = std::mem::take(self.join_set.lock().unwrap().deref_mut()); - for res in join_set.join_all().await { - res?; + match tokio::time::timeout_at(deadline, join_set.join_all()).await { + Ok(results) => { + for res in results { + res?; + } + } + // Dropping the JoinSet aborts the stalled forwarding tasks. Their receivers + // observe EOS and mark missing reports as terminal, rather than hanging queries. + Err(_) => { + log::warn!("timed out draining worker coordinator channels after query completion") + } } Ok(()) } @@ -347,6 +366,16 @@ impl<'a> StageCoordinator<'a> { } } } + // The worker->coordinator stream has ended, so no further reports will arrive for + // this task. If the task never reported (e.g. it was planned but never executed + // because execution short-circuited, or the stream ended abruptly), mark it terminal + // so that waiters do not block forever. Reported values are left untouched. + if let Some(task_metrics) = &task_metrics { + task_metrics.mark_terminal(task_key); + } + if let Some(store) = &completed_dynamic_filter_store { + store.mark_terminal(task_key); + } }); load_info_rx } @@ -514,11 +543,17 @@ struct TaskSpecializedPlan { dynamic_filter_remote_producer_ids: Vec, } -pub(super) struct NotifyGuard(Arc); +pub(super) struct NotifyGuard { + notifier: Arc, + deadline: watch::Sender>, + timeout: Duration, +} impl Drop for NotifyGuard { fn drop(&mut self) { - self.0.notify_waiters(); + self.deadline + .send_replace(Some(TokioInstant::now() + self.timeout)); + self.notifier.notify_waiters(); } } @@ -559,6 +594,9 @@ impl CoordinatorToWorkerMetrics { #[cfg(test)] mod tests { use super::*; + use datafusion::arrow::datatypes::Schema; + use datafusion::physical_plan::empty::EmptyExec; + use std::time::Duration; /// Regression test for a rustc miscompilation (present at least through /// 1.96, fixed in 1.98) of [`CoordinatorToWorkerMetrics::new`]. @@ -625,4 +663,96 @@ mod tests { ); } } + + #[tokio::test] + async fn channel_close_without_report_is_terminal_but_incomplete() { + let task_ctx = Arc::new(TaskContext::default()); + let metrics = ExecutionPlanMetricsSet::new(); + let metrics_store = Arc::new(Store::::new()); + let filters_store = Arc::new(Store::::new()); + let coordinator = QueryCoordinator::new( + task_ctx, + &metrics, + Some(Arc::clone(&metrics_store)), + Some(Arc::clone(&filters_store)), + ); + let stage = LocalStage { + query_id: Uuid::new_v4(), + num: 1, + plan: Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + tasks: 1, + metrics_set: Default::default(), + }; + let mut stage_coordinator = coordinator.stage_coordinator(&stage); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let _load_info = stage_coordinator.worker_to_coordinator_task(0, rx); + let key = TaskKey { + query_id: stage.query_id, + stage_id: stage.num, + task_number: 0, + }; + // Keep the channel open: the task must still be pending, not prematurely terminal. + assert_eq!(metrics_store.snapshot(&[key]).pending, vec![key]); + drop(tx); + let reported = tokio::time::timeout( + Duration::from_secs(1), + metrics_store.wait_for_terminal(&[key]), + ) + .await + .unwrap(); + assert_eq!(reported.terminal_without_report, vec![key]); + assert!(reported.reported.is_empty()); + assert_eq!( + filters_store + .wait_for_terminal(&[key]) + .await + .terminal_without_report, + vec![key] + ); + } + + #[tokio::test] + async fn stuck_channel_drain_is_bounded_and_does_not_fail_a_finished_query() { + let metrics = ExecutionPlanMetricsSet::new(); + let metrics_store = Arc::new(Store::::new()); + let coordinator = Arc::new(QueryCoordinator::new( + Arc::new(TaskContext::default()), + &metrics, + Some(Arc::clone(&metrics_store)), + None, + )); + let stage = LocalStage { + query_id: Uuid::new_v4(), + num: 1, + plan: Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + tasks: 1, + metrics_set: Default::default(), + }; + let key = TaskKey { + query_id: stage.query_id, + stage_id: stage.num, + task_number: 0, + }; + // Hold a live sender inside the stuck forwarding task. + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + coordinator + .stage_coordinator(&stage) + .worker_to_coordinator_task(0, rx); + coordinator.join_set.lock().unwrap().spawn(async move { + let _tx = tx; + futures::future::pending::<()>().await; + Ok(()) + }); + coordinator + .drain_pending_tasks(TokioInstant::now() + Duration::from_millis(25)) + .await + .unwrap(); + let snapshot = tokio::time::timeout( + Duration::from_secs(1), + metrics_store.wait_for_terminal(&[key]), + ) + .await + .unwrap(); + assert_eq!(snapshot.terminal_without_report, vec![key]); + } } diff --git a/src/coordinator/store.rs b/src/coordinator/store.rs index 806240107..de87fa6e7 100644 --- a/src/coordinator/store.rs +++ b/src/coordinator/store.rs @@ -6,7 +6,25 @@ use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; use tokio::sync::watch; -type StoreMap = HashMap; +/// Terminal state of a task-scoped entry in a [Store]. Tasks without an entry are still pending. +#[derive(Debug, Clone)] +enum TaskEntry { + /// The task reported a value. + Reported(T), + /// The task finished without reporting a value, for example because it was planned but never + /// executed, or because its worker->coordinator stream ended abruptly. + TerminalWithoutReport, +} + +type StoreMap = HashMap>; + +/// Point-in-time task state for the expected tasks in a plan. +#[derive(Debug)] +pub(crate) struct StoreSnapshot { + pub(crate) reported: HashMap, + pub(crate) pending: Vec, + pub(crate) terminal_without_report: Vec, +} /// Stores task-scoped values and notifies waiters when entries change. #[derive(Debug, Clone)] @@ -21,13 +39,35 @@ impl Store { Self { tx, rx } } + /// Records the value reported by a task, making it terminal. pub(crate) fn insert(&self, key: TaskKey, value: T) { self.tx.send_modify(|map| { - map.insert(key, value); + map.insert(key, TaskEntry::Reported(value)); + }); + } + + /// Marks a task as terminal without a reported value. Does nothing if the task already + /// reported a value. + pub(crate) fn mark_terminal(&self, key: TaskKey) { + self.tx.send_if_modified(|map| { + if map.contains_key(&key) { + return false; + } + map.insert(key, TaskEntry::TerminalWithoutReport); + true }); } - pub(crate) async fn wait_for(&self, expected_keys: &[TaskKey]) -> StoreMap + /// Waits until every expected task is terminal, and returns the values of the tasks that + /// reported one. Tasks that became terminal without reporting are absent from the result. + pub(crate) async fn wait_for(&self, expected_keys: &[TaskKey]) -> HashMap + where + T: Clone, + { + self.wait_for_terminal(expected_keys).await.reported + } + + pub(crate) async fn wait_for_terminal(&self, expected_keys: &[TaskKey]) -> StoreSnapshot where T: Clone, { @@ -37,8 +77,39 @@ impl Store { .wait_for(|map| expected_keys.iter().all(|key| map.contains_key(key))) .await; } - rx.borrow().clone() + snapshot(&rx.borrow(), expected_keys) } + + pub(crate) fn snapshot(&self, expected_keys: &[TaskKey]) -> StoreSnapshot + where + T: Clone, + { + snapshot(&self.rx.borrow(), expected_keys) + } +} + +fn snapshot(map: &StoreMap, expected_keys: &[TaskKey]) -> StoreSnapshot { + let mut result = StoreSnapshot { + // Keep every available report, even if it is not in the prepared plan's expected keys. + // The expected keys determine completeness, not which reports are observable. + reported: map + .iter() + .filter_map(|(key, entry)| match entry { + TaskEntry::Reported(value) => Some((*key, value.clone())), + TaskEntry::TerminalWithoutReport => None, + }) + .collect(), + pending: vec![], + terminal_without_report: vec![], + }; + for key in expected_keys { + match map.get(key) { + Some(TaskEntry::Reported(_)) => {} + Some(TaskEntry::TerminalWithoutReport) => result.terminal_without_report.push(*key), + None => result.pending.push(*key), + } + } + result } pub(crate) fn task_keys_for_plan(plan: &Arc) -> Vec { @@ -58,3 +129,58 @@ pub(crate) fn task_keys_for_plan(plan: &Arc) -> Vec }); task_keys } + +#[cfg(test)] +mod tests { + use super::*; + use uuid::Uuid; + + fn task_key(query_id: Uuid, task_number: usize) -> TaskKey { + TaskKey { + query_id, + stage_id: 1, + task_number, + } + } + + #[test] + fn mark_terminal_does_not_override_reported_value() { + let query_id = Uuid::new_v4(); + let key = task_key(query_id, 0); + let store = Store::new(); + store.insert(key, 10); + store.mark_terminal(key); + + let snapshot = store.snapshot(&[key]); + assert_eq!(snapshot.reported.get(&key), Some(&10)); + assert!(snapshot.pending.is_empty()); + assert!(snapshot.terminal_without_report.is_empty()); + } + + #[test] + fn snapshot_is_non_blocking_and_distinguishes_pending_from_lost() { + let query_id = Uuid::new_v4(); + let keys = [ + task_key(query_id, 0), + task_key(query_id, 1), + task_key(query_id, 2), + ]; + let store = Store::new(); + store.insert(keys[2], 12); + store.mark_terminal(keys[0]); + + let partial = store.snapshot(&keys); + assert_eq!(partial.reported, HashMap::from([(keys[2], 12)])); + assert_eq!(partial.pending, vec![keys[1]]); + assert_eq!(partial.terminal_without_report, vec![keys[0]]); + + store.insert(keys[1], 11); + let unplanned = task_key(query_id, 10); + store.insert(unplanned, 110); + let terminal = store.snapshot(&keys); + assert!(terminal.pending.is_empty()); + assert_eq!(terminal.reported.len(), 3); + assert_eq!(terminal.reported.get(&unplanned), Some(&110)); + assert_eq!(terminal.terminal_without_report, vec![keys[0]]); + } +} diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index 4b8fd845e..99c22b9ec 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -4,6 +4,8 @@ use datafusion::config::{ConfigExtension, ConfigOptions}; use datafusion::execution::TaskContext; use datafusion::prelude::SessionConfig; +pub(crate) const DEFAULT_METRICS_FINALIZATION_TIMEOUT_MS: u64 = 5_000; + extensions_options! { /// Configuration for the distributed planner. pub struct DistributedConfig { @@ -29,6 +31,10 @@ extensions_options! { /// Propagate collected metrics from all nodes in the plan across network boundaries /// so that they can be reconstructed on the head node of the plan. pub collect_metrics: bool, default = true + /// Milliseconds allowed after the result stream ends for coordinator channels to close + /// and worker reports (metrics and completed dynamic filters) to arrive. All finalization + /// steps share this one deadline. A value of 0 does not wait for late reports. + pub metrics_finalization_timeout_ms: u64, default = DEFAULT_METRICS_FINALIZATION_TIMEOUT_MS /// Collect completed dynamic filters from worker tasks so that they can be displayed in /// the distributed plan. This does not control whether dynamic filtering is used during /// query execution. diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index 273d64d0d..f889f094e 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/src/distributed_planner/distributed_query_planner.rs @@ -24,6 +24,7 @@ use datafusion::physical_plan::{ use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; use futures::future::BoxFuture; use std::sync::Arc; +use std::time::Duration; /// Transforms a single-node physical plan into a distributed plan by injecting network /// boundaries between stages. @@ -174,7 +175,10 @@ fn create_distributed_exec( Arc::new( DistributedExec::new(plan) .with_metrics_collection(d_cfg.collect_metrics) - .with_dynamic_filter_collection(d_cfg.collect_dynamic_filters), + .with_dynamic_filter_collection(d_cfg.collect_dynamic_filters) + .with_finalization_timeout(Duration::from_millis( + d_cfg.metrics_finalization_timeout_ms, + )), ) } diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index de006f1b6..2ad7f8b42 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -11,6 +11,7 @@ mod push_fetch_into_network_coalesce; mod session_state_builder_ext; mod statistics; +pub(crate) use distributed_config::DEFAULT_METRICS_FINALIZATION_TIMEOUT_MS; pub use distributed_config::DistributedConfig; pub(crate) use inject_network_boundaries::{ InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, inject_network_boundaries, diff --git a/src/execution_plans/sampler.rs b/src/execution_plans/sampler.rs index d393efb26..f5d42628a 100644 --- a/src/execution_plans/sampler.rs +++ b/src/execution_plans/sampler.rs @@ -110,6 +110,10 @@ pub(crate) struct SamplerGate { } impl SamplerGate { + pub(crate) fn has_sampler(&self) -> bool { + self.sampling_start_tx_opt.is_some() + } + /// Takes the [LoadInfo] receivers. This function can only be called once, as subsequent calls /// will always return an empty [Vec]. pub(crate) fn take_receivers(&mut self) -> Vec> { @@ -146,6 +150,7 @@ impl SamplerExec { partition_idx: i, input: Arc::clone(&input), stream: Mutex::new(None), + sampling_done_rx: Mutex::new(None), metrics: Arc::clone(&metrics), kick_off_at: Arc::new(OnceLock::new()), first_batch_at: Arc::new(OnceLock::new()), @@ -186,12 +191,37 @@ impl SamplerExec { receivers, }) } + + /// Stop samplers whose task never received ExecuteTask, then wait for their futures to be + /// dropped before reading the plan's accounting metrics. Otherwise in-flight scans can + /// continue incrementing I/O counters after their snapshot has been sent. + pub(crate) async fn stop_unexecuted_sampling(plan: &Arc) -> Result<()> { + let mut done = vec![]; + plan.apply(|node| { + if let Some(sampler) = node.downcast_ref::() { + for partition in &sampler.partition_samplers { + // The saved stream owns the SpawnedTask; dropping it cancels any work not + // consumed by an ExecuteTask request. + drop(partition.stream.lock().unwrap().take()); + if let Some(rx) = partition.sampling_done_rx.lock().unwrap().take() { + done.push(rx); + } + } + } + Ok(TreeNodeRecursion::Continue) + })?; + for rx in done { + let _ = rx.await; + } + Ok(()) + } } pub(crate) struct PartitionSampler { partition_idx: usize, input: Arc, stream: Mutex>, + sampling_done_rx: Mutex>>, execution_started: Arc, // Metrics state. @@ -282,7 +312,12 @@ impl PartitionSampler { first_batch_at: Arc::clone(&self.first_batch_at), }; + let (sampling_done_tx, sampling_done_rx) = oneshot::channel::<()>(); + let task = SpawnedTask::spawn(async move { + // On completion or abort, dropping this sender signals that the sampler has + // stopped touching the input plan's metrics. + let _sampling_done_tx = sampling_done_tx; sampling_start_rx .wait_for(|started| *started) .await @@ -333,6 +368,10 @@ impl PartitionSampler { .lock() .expect("poisoned lock") .replace(Box::pin(RecordBatchStreamAdapter::new(schema, stream))); + self.sampling_done_rx + .lock() + .unwrap() + .replace(sampling_done_rx); Ok(sampling_rx) } diff --git a/src/lib.rs b/src/lib.rs index 523252982..f345502f4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,7 +19,7 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +pub use coordinator::{DistributedExec, TaskMetricsSnapshot}; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, ProducerHead, SessionStateBuilderExt, diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs index fd1139426..32a9863ab 100644 --- a/src/metrics/mod.rs +++ b/src/metrics/mod.rs @@ -15,3 +15,7 @@ pub use task_metrics_rewriter::{DistributedMetricsFormat, rewrite_distributed_pl /// Label used to annotate metrics in execution plan nodes with the task in which they were executed. /// Note that the same task id may be used in multiple stages. pub const DISTRIBUTED_DATAFUSION_TASK_ID_LABEL: &str = "task_id"; + +/// Task-level marker for a plan that was installed (and possibly sampled) but not executed +/// through ExecuteTask. This uses the existing metrics wire format. +pub(crate) const TASK_NOT_EXECUTED_METRIC: &str = "task_not_executed"; diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 4c45aae72..4061fc37a 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -40,7 +40,9 @@ impl DistributedMetricsFormat { /// Rewrites a distributed plan with metrics. Does nothing if the root node is not a [DistributedExec]. /// Returns an error if the distributed plan was not executed. /// -/// Waits for all worker task metrics to arrive before rewriting, so the result is always complete. +/// Waits for every worker task to report (including sampled and skipped tasks) before rewriting. +/// Returns an error rather than silently discarding accounting metrics if a channel closes +/// without a report or the finalization timeout elapses. pub async fn rewrite_distributed_plan_with_metrics( plan: Arc, format: DistributedMetricsFormat, @@ -54,9 +56,7 @@ pub async fn rewrite_distributed_plan_with_metrics( } let head_stage = distributed_exec.head_stage()?; - let Some(metrics_collection) = distributed_exec.wait_for_metrics().await else { - return internal_err!("metrics were enabled but the execution was not prepared"); - }; + let metrics_collection = distributed_exec.complete_metrics().await?; let task_metrics = collect_plan_metrics(&head_stage)?; // Rewrite the DistributedExec's child plan with metrics. @@ -110,8 +110,7 @@ fn stage_metrics( }; let Some(task_metrics) = metrics_collection.get(&task_key) else { return internal_err!( - "not enough metrics provided to rewrite task: missing metrics for task {} in stage {}", - task_number, + "missing metrics report for task {task_number} in stage {}", stage.num ); }; @@ -271,12 +270,18 @@ pub fn stage_metrics_rewriter( }; let Some(task_metrics) = metrics_collection.get(&task_key) else { return internal_err!( - "not enough metrics provided to rewrite task: missing metrics for task {} in stage {}", - task_id, + "missing metrics report for task {task_id} in stage {}", stage.num ); }; + // A statically planned task may never receive ExecuteTask. Its lifecycle was reported, + // but there are no node execution metrics to attach. AQE tasks that sampled data do + // provide the actual sampled node metrics and are rewritten below. + if task_metrics.was_not_executed() && task_metrics.pre_order_plan_metrics.is_empty() { + continue; + } + let mut per_task_counter = 0usize; stage.plan.apply_with_dt_ctx(d_ctx, |node, _ctx| { if per_task_counter >= task_metrics.pre_order_plan_metrics.len() { @@ -347,6 +352,7 @@ mod tests { use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; use itertools::Itertools; + use std::borrow::Cow; use std::sync::Arc; use test_case::test_case; use uuid::Uuid; @@ -592,6 +598,82 @@ mod tests { .await; } + /// A skipped task has a lifecycle report but contributes no fabricated node metrics. + #[tokio::test] + async fn test_stage_metrics_rewriter_with_skipped_task() { + let ctx = make_test_ctx().await; + let plan = ctx + .sql("SELECT id, COUNT(*) as count FROM table1 WHERE id > 1 GROUP BY id") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let mut stage = make_test_stage(plan.clone()); + stage.tasks = 2; + let num_nodes = count_plan_nodes_up_to_network_boundary(&plan); + + let task_key = |task_number| TaskKey { + query_id: stage.query_id, + stage_id: stage.num, + task_number, + }; + let count = Count::new(); + count.add(1); + let mut skipped_metrics = MetricsSet::new(); + skipped_metrics.push(Arc::new(Metric::new( + MetricValue::Count { + name: Cow::Borrowed(crate::metrics::TASK_NOT_EXECUTED_METRIC), + count, + }, + None, + ))); + let metrics_collection = HashMap::from([ + ( + task_key(0), + TaskMetrics { + task_metrics: skipped_metrics, + pre_order_plan_metrics: vec![], + }, + ), + ( + task_key(1), + TaskMetrics { + task_metrics: MetricsSet::new(), + pre_order_plan_metrics: (0..num_nodes) + .map(|_| make_test_metrics_set_from_seed(1, 1)) + .collect(), + }, + ), + ]); + + let stage_metrics = super::stage_metrics(&stage, &metrics_collection).unwrap(); + assert!( + stage_metrics + .iter() + .any(|metric| metric.value().name() == crate::metrics::TASK_NOT_EXECUTED_METRIC) + ); + + let rewritten_plan = stage_metrics_rewriter( + &stage, + &metrics_collection, + DistributedMetricsFormat::PerTask, + ) + .unwrap(); + let mut actual_metrics = vec![]; + collect_metrics_from_plan(&rewritten_plan, &mut actual_metrics); + assert_eq!(actual_metrics.len(), num_nodes); + for node_metrics in actual_metrics { + let task_ids: Vec<_> = node_metrics + .iter() + .flat_map(|m| m.labels().iter()) + .filter(|l| l.name() == DISTRIBUTED_DATAFUSION_TASK_ID_LABEL) + .map(|l| l.value().to_string()) + .collect(); + assert_eq!(task_ids, vec!["1"]); + } + } + #[tokio::test] async fn test_rewrite_unexecuted_distributed_plan_with_metrics_err() { let ctx = make_test_distributed_ctx().await; diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 4f5257cfc..77ae75d79 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -176,6 +176,16 @@ pub struct TaskMetrics { pub task_metrics: MetricsSet, } +impl TaskMetrics { + /// Whether this report came from a planned task that never received ExecuteTask. Its plan + /// may still contain real sampling metrics, so it must not be treated as a zero-cost task. + pub fn was_not_executed(&self) -> bool { + self.task_metrics + .iter() + .any(|metric| metric.value().name() == crate::metrics::TASK_NOT_EXECUTED_METRIC) + } +} + #[derive(Default)] pub struct LoadInfo { /// The partition index to which this message belongs to. diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index bef27f3b9..b5403f4a8 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -162,7 +162,17 @@ impl Worker { #[allow(clippy::disallowed_methods)] tokio::spawn(async move { let mut stream = stream.map_ok(set_work_unit_received_time); - while let Some(Ok(msg)) = stream.next().await { + let mut sampling_started = false; + let mut stream_failed = false; + while let Some(msg) = stream.next().await { + let msg = match msg { + Ok(msg) => msg, + Err(err) => { + log::warn!("coordinator channel closed with an error: {err}"); + stream_failed = true; + break; + } + }; match msg { CoordinatorToWorkerMsg::WorkUnitBatch(work_unit_batch) => { let Some(work_unit_senders) = work_unit_senders.as_mut() else { @@ -193,6 +203,7 @@ impl Worker { let _ = work_unit_senders.take(); } CoordinatorToWorkerMsg::KickOffSampling => { + sampling_started = sampler_gate.has_sampler(); sampler_gate.kick_off(); } } @@ -211,22 +222,48 @@ impl Worker { .lock() .unwrap() .take(); - if let Some(Ok(plan)) = task_data.final_plan.get() { - let d_ctx = DistributedTaskContext { - task_index: key.task_number, - task_count: request.task_count, + if stream_failed { + // A transport failure is not a clean end-of-query signal. Do not send a + // "complete" report from an in-flight plan; the coordinator will observe + // terminal-without-report and surface incomplete accounting. + drop(metrics_tx); + drop(dynamic_filters_tx); + task_data_entries.invalidate(&key).await; + return; + } + let d_ctx = DistributedTaskContext { + task_index: key.task_number, + task_count: request.task_count, + }; + let final_plan = task_data + .final_plan + .get() + .and_then(|result| result.as_ref().ok()); + if final_plan.is_none() { + // Stabilize sampled scan counters before sending a report from the base plan. + let _ = SamplerExec::stop_unexecuted_sampling(&task_data.base_plan).await; + } else { + task_data.task_data_metrics.mark_execution_finished(); + } + if let Some(metrics_tx) = metrics_tx { + let (plan, task_metrics) = match final_plan { + Some(plan) => (Some(plan), task_data.task_data_metrics.to_metrics_set()), + None => { + // An AQE task can sample without executing. The base plan has its I/O + // metrics; a statically skipped task only reports its lifecycle. + ( + sampling_started.then_some(&task_data.base_plan), + task_data.task_data_metrics.to_unexecuted_metrics_set(), + ) + } }; - 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); - } - if let Some(dynamic_filters_tx) = dynamic_filters_tx { - // TODO(#686): handle error - let dynamic_filters = - build_task_completed_dynamic_filters(plan).unwrap_or_default(); - let _ = dynamic_filters_tx.send(dynamic_filters); - } + send_metrics_via_channel(metrics_tx, plan, d_ctx, task_metrics); + } + if let (Some(plan), Some(dynamic_filters_tx)) = (final_plan, dynamic_filters_tx) { + // TODO(#686): handle error + let dynamic_filters = + build_task_completed_dynamic_filters(plan).unwrap_or_default(); + let _ = dynamic_filters_tx.send(dynamic_filters); } task_data_entries.invalidate(&key).await }); @@ -358,19 +395,97 @@ fn build_task_completed_dynamic_filters( /// coordinator channel oneshot. fn send_metrics_via_channel( metrics_tx: Sender, - plan: &Arc, + plan: Option<&Arc>, dt_ctx: DistributedTaskContext, - task_data_metrics: &Arc, + task_metrics: datafusion::physical_plan::metrics::MetricsSet, ) { let mut pre_order_plan_metrics = vec![]; - let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { - pre_order_plan_metrics.push(node.metrics().unwrap_or_default()); - Ok(TreeNodeRecursion::Continue) - }); + if let Some(plan) = plan { + let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { + pre_order_plan_metrics.push(node.metrics().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: task_data_metrics.to_metrics_set(), + task_metrics, }); } + +#[cfg(test)] +mod tests { + use super::*; + use crate::worker::task_data::PLAN_EXECUTED_AT_METRIC; + use datafusion::arrow::datatypes::Schema; + use datafusion::physical_plan::empty::EmptyExec; + use uuid::Uuid; + + fn request() -> SetPlanRequest { + SetPlanRequest { + task_key: crate::TaskKey { + query_id: Uuid::new_v4(), + stage_id: 1, + task_number: 0, + }, + task_count: 1, + plan: MaybeEncoded::Decoded(Arc::new(EmptyExec::new(Arc::new(Schema::empty())))), + dynamic_filter_remote_producer_ids: vec![], + work_unit_feed_declarations: vec![], + target_worker_url: "http://localhost:8080".parse().unwrap(), + query_start_time_ns: 0, + } + } + + #[tokio::test] + async fn channel_eos_reports_unexecuted_task_without_execution_metrics() { + let worker = Worker::default(); + let result = worker + .coordinator_channel( + HeaderMap::new(), + request(), + futures::stream::empty().boxed(), + ) + .await + .unwrap(); + let messages = result.stream.try_collect::>().await.unwrap(); + let reports: Vec<_> = messages + .into_iter() + .filter_map(|msg| match msg { + WorkerToCoordinatorMsg::TaskMetrics(report) => Some(report), + _ => None, + }) + .collect(); + assert_eq!(reports.len(), 1); + assert!(reports[0].was_not_executed()); + assert!(reports[0].pre_order_plan_metrics.is_empty()); + assert!( + !reports[0] + .task_metrics + .iter() + .any(|metric| metric.value().name() == PLAN_EXECUTED_AT_METRIC) + ); + } + + #[tokio::test] + async fn inbound_channel_error_does_not_report_partial_metrics_as_complete() { + let result = Worker::default() + .coordinator_channel( + HeaderMap::new(), + request(), + futures::stream::once(async { + Err(exec_datafusion_err!("injected transport failure")) + }) + .boxed(), + ) + .await + .unwrap(); + let messages = result.stream.try_collect::>().await.unwrap(); + assert!( + !messages + .iter() + .any(|msg| matches!(msg, WorkerToCoordinatorMsg::TaskMetrics(_))) + ); + } +} diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index 81ad44617..aad2f66f3 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -1,10 +1,11 @@ use crate::common::OnceLockResult; use crate::common::now_ns; +use crate::metrics::TASK_NOT_EXECUTED_METRIC; use crate::{MaxLatencyMetric, ProducerHead, TaskCompletedDynamicFilters, TaskMetrics}; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_plan::metrics::{Metric, MetricValue, MetricsSet}; +use datafusion::physical_plan::metrics::{Count, Metric, MetricValue, MetricsSet}; use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; @@ -92,6 +93,26 @@ impl TaskDataMetrics { metrics_set } + + /// Planning and sampling may have happened, but ExecuteTask was never called. Do not + /// manufacture timestamps for execution or completion. + pub(super) fn to_unexecuted_metrics_set(&self) -> MetricsSet { + let mut metrics_set = MetricsSet::new(); + metrics_set.push(max_latency_metric( + PLAN_ADDED_AT_METRIC, + &self.plan_added_at, + )); + let count = Count::new(); + count.add(1); + metrics_set.push(Arc::new(Metric::new( + MetricValue::Count { + name: Cow::Borrowed(TASK_NOT_EXECUTED_METRIC), + count, + }, + None, + ))); + metrics_set + } } fn max_latency_metric(name: &'static str, value: &MaxLatencyMetric) -> Arc { diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index c6a421da7..724e5647e 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -9,6 +9,7 @@ mod tests { use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ExecutionPlan, execute_stream}; use datafusion::prelude::SessionContext; + use datafusion_distributed::DistributedExec; use datafusion_distributed::test_utils::localhost::start_localhost_context; use datafusion_distributed::test_utils::parquet::register_parquet_tables; use datafusion_distributed::test_utils::test_work_unit_feed::{ @@ -22,6 +23,7 @@ mod tests { }; use futures::TryStreamExt; use std::sync::Arc; + use std::time::Duration; use test_case::test_case; #[test_case(DistributedMetricsFormat::Aggregated ; "aggregated_metrics")] @@ -379,6 +381,120 @@ mod tests { Ok(()) } + /// Regression test for https://github.com/datafusion-contrib/datafusion-distributed/issues/739 + /// + /// In a partitioned hash join whose build side is empty, DataFusion completes the join + /// without ever polling the probe side. The probe-side worker tasks are therefore planned + /// (they receive the plan through the coordinator channel) but never executed, so they never + /// produce a `TaskMetrics` report. Metrics finalization must not wait for them forever. + #[tokio::test] + async fn test_metrics_collection_with_unexecuted_probe_side_tasks() + -> Result<(), Box> { + let (mut d_ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; + d_ctx.set_distributed_dynamic_task_count(true)?; + register_parquet_tables(&d_ctx).await?; + force_partitioned_hash_joins(&d_ctx); + + let query = r#" + SELECT a."MinTemp", b."MaxTemp" + FROM weather a + JOIN weather b ON a."RainToday" = b."RainToday" + WHERE a."MinTemp" > 1000000 + "#; + + let plan = d_ctx.sql(query).await?.create_physical_plan().await?; + let batches = execute_stream(plan.clone(), d_ctx.task_ctx())? + .try_collect::>() + .await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + + let with_filters = tokio::time::timeout( + Duration::from_secs(10), + datafusion_distributed::rewrite_distributed_plan_with_dynamic_filters( + plan.clone(), + &d_ctx.task_ctx(), + ), + ) + .await??; + rewrite_with_metrics_within_timeout(with_filters).await?; + let snapshot = plan + .downcast_ref::() + .unwrap() + .metrics_snapshot() + .unwrap(); + assert!( + snapshot.all_reported(), + "unexpected lost report: {:?}", + snapshot.missing_reports + ); + let unexecuted = snapshot.unexecuted_tasks(); + assert!(!unexecuted.is_empty(), "AQE did not skip any task"); + assert!( + unexecuted.iter().any(|key| { + snapshot.reported[key] + .pre_order_plan_metrics + .iter() + .flat_map(|set| set.iter()) + .any(|metric| { + metric.value().name() == "files_opened" && metric.value().as_usize() > 0 + }) + }), + "skipped AQE tasks lost their actual sampling I/O metrics" + ); + Ok(()) + } + + /// Regression test for https://github.com/datafusion-contrib/datafusion-distributed/issues/739 + /// + /// A static multi-task scan with `LIMIT 1` can be satisfied before every planned task + /// executes. Metrics finalization must still complete. + /// + /// Whether a task is skipped is timing dependent; the worker-channel unit test covers that + /// case deterministically, while this test exercises the public rewrite after a static LIMIT. + #[tokio::test] + async fn test_metrics_collection_with_limit_short_circuiting_static_tasks() + -> Result<(), Box> { + let (mut d_ctx, _guard, _) = start_localhost_context(8, DefaultSessionBuilder).await; + d_ctx.set_distributed_file_scan_config_bytes_per_partition(1024 * 1024)?; + register_parquet_tables(&d_ctx).await?; + + let plan = d_ctx + .sql(r#"SELECT "FL_DATE" FROM flights_1m LIMIT 1"#) + .await? + .create_physical_plan() + .await?; + assert!(plan.is::(), "expected a distributed scan"); + let batches = execute_stream(plan.clone(), d_ctx.task_ctx())? + .try_collect::>() + .await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); + rewrite_with_metrics_within_timeout(plan).await?; + Ok(()) + } + + /// Rewrites the plan with metrics, failing if it does not finish within a generous timeout + /// instead of hanging the test suite. + async fn rewrite_with_metrics_within_timeout( + plan: Arc, + ) -> Result, Box> { + let rewrite = + rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::PerTask); + match tokio::time::timeout(Duration::from_secs(10), rewrite).await { + Ok(result) => Ok(result?), + Err(_) => { + Err("metrics finalization did not complete: a task never became terminal".into()) + } + } + } + + fn force_partitioned_hash_joins(ctx: &SessionContext) { + let session_state = ctx.state_ref(); + let mut session_state = session_state.write(); + let options = session_state.config_mut().options_mut(); + options.optimizer.hash_join_single_partition_threshold = 0; + options.optimizer.hash_join_single_partition_threshold_rows = 0; + } + /// Looks for an [ExecutionPlan] that matches the provided type parameter `T1` in /// the left node and `T2` in the right node and compares its metrics. /// There might be more than one, so `index` determines which one is compared. From 638fd0bb1ab37a6f9b7764697f24c6cad3e0b516 Mon Sep 17 00:00:00 2001 From: Sergio Esteves Date: Fri, 25 Sep 2026 10:34:59 +0100 Subject: [PATCH 2/2] refactor: defer partial metrics snapshot API --- docs/source/user-guide/05-metrics.md | 29 ------------- src/coordinator/distributed.rs | 63 +++------------------------- src/coordinator/mod.rs | 2 +- src/lib.rs | 2 +- src/protocol/worker_channel.rs | 2 +- tests/metrics_collection.rs | 29 ++----------- 6 files changed, 12 insertions(+), 115 deletions(-) diff --git a/docs/source/user-guide/05-metrics.md b/docs/source/user-guide/05-metrics.md index 062b06cad..3d27bca23 100644 --- a/docs/source/user-guide/05-metrics.md +++ b/docs/source/user-guide/05-metrics.md @@ -107,32 +107,3 @@ runtime metrics, including network-level metrics on the boundaries: > If `plan` is not a distributed plan (its root is not a `DistributedExec`), > `rewrite_distributed_plan_with_metrics` returns it unchanged, so it is always safe to call. - -## Inspecting a partial snapshot - -`DistributedExec::metrics_snapshot()` returns the metrics received **so far** without waiting. -It returns `None` if collection is disabled or the plan has not yet been prepared. The snapshot -has separate `pending`, `missing_reports`, and `reported` task keys; `all_terminal()` says whether -every expected task has reported or its stream has ended. `all_reported()` says whether metrics -are available for every task; use it when accounting must be complete. `unexecuted_tasks()` identifies planned -tasks that reported actual AQE sampling metrics but never received an `ExecuteTask` call. These -tasks have a `task_not_executed=1` stage metric instead of fabricated execution timestamps. - -```rust -use datafusion_distributed::DistributedExec; - -if let Some(exec) = plan.downcast_ref::() { - if let Some(snapshot) = exec.metrics_snapshot() { - println!("{} reports; {} pending", snapshot.reported.len(), snapshot.pending.len()); - if !snapshot.all_reported() { - // Do not treat a missing task report as zero I/O. - println!("missing reports: {:?}", snapshot.missing_reports); - } - } -} -``` - -The existing `DistributedExec::wait_for_metrics()` only returns `Some` for a complete set of -reports. If any task fails to report or the wait times out, it returns `None`; call -`metrics_snapshot()` to inspect what arrived. A closed stream without a metrics report is not -equivalent to an unexecuted task with a report. diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index ecb34f26b..ac056415a 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -2,7 +2,7 @@ use crate::common::require_one_child; use crate::coordinator::prepare_dynamic_plan::prepare_dynamic_plan; use crate::coordinator::prepare_static_plan::prepare_static_plan; use crate::coordinator::query_coordinator::QueryCoordinator; -use crate::coordinator::store::{Store, StoreSnapshot, task_keys_for_plan}; +use crate::coordinator::store::{Store, task_keys_for_plan}; use crate::distributed_planner::DEFAULT_METRICS_FINALIZATION_TIMEOUT_MS; use crate::dynamic_filtering::{ is_dynamic_filtering_enabled, sever_dynamic_filter_relationships_in_plan_for_display, @@ -24,46 +24,6 @@ use std::time::Duration; use tokio::sync::watch; use tokio::time::Instant; -/// Non-blocking, point-in-time view of the reports expected from a prepared distributed plan. -#[derive(Debug)] -pub struct TaskMetricsSnapshot { - /// Metrics received so far, including actual sampling metrics from tasks that never ran. - pub reported: HashMap, - /// Tasks whose coordinator channels are still open and have not reported metrics. - pub pending: Vec, - /// Tasks whose channels closed without a report. Their work cannot be accounted for. - pub missing_reports: Vec, -} - -impl TaskMetricsSnapshot { - fn from_store(snapshot: StoreSnapshot) -> Self { - Self { - reported: snapshot.reported, - pending: snapshot.pending, - missing_reports: snapshot.terminal_without_report, - } - } - - /// Whether every expected task has either reported or had its stream end without a report. - pub fn all_terminal(&self) -> bool { - self.pending.is_empty() - } - - /// Whether every expected task has actually reported metrics. This distinguishes a - /// finished-but-incomplete query from one whose measurements are complete. - pub fn all_reported(&self) -> bool { - self.all_terminal() && self.missing_reports.is_empty() - } - - /// Tasks that reported real sampling metrics but never received ExecuteTask. - pub fn unexecuted_tasks(&self) -> Vec { - self.reported - .iter() - .filter_map(|(key, report)| report.was_not_executed().then_some(*key)) - .collect() - } -} - /// [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. Unless explicitly set in @@ -149,23 +109,11 @@ impl DistributedExec { } /// Waits for complete metrics, if collection is enabled and execution has been prepared. - /// Returns `None` if a task failed to report or the finalization timeout elapsed; use - /// [`Self::metrics_snapshot`] to inspect reports and their completeness in that case. + /// Returns `None` if a task failed to report or the finalization timeout elapsed. pub async fn wait_for_metrics(&self) -> Option> { self.complete_metrics().await.ok() } - /// Returns all available metrics immediately, including gaps in task numbers, along with - /// pending and terminal-without-report task keys. Returns `None` if metrics are disabled or - /// the distributed plan has not yet been prepared. - pub fn metrics_snapshot(&self) -> Option { - let store = self.metrics_store.as_ref()?; - let plan = &self.prepared_plan.get()?.plan_for_viz; - Some(TaskMetricsSnapshot::from_store( - store.snapshot(&task_keys_for_plan(plan)), - )) - } - pub(crate) async fn complete_metrics(&self) -> Result> { let store = self .metrics_store @@ -474,10 +422,9 @@ mod tests { ) .await; assert!(result.unwrap_err().to_string().contains("metrics missing")); - let snapshot = TaskMetricsSnapshot::from_store(store.snapshot(&[reported, lost])); - assert!(snapshot.all_terminal()); - assert!(!snapshot.all_reported()); - assert_eq!(snapshot.missing_reports, vec![lost]); + let snapshot = store.snapshot(&[reported, lost]); + assert!(snapshot.pending.is_empty()); + assert_eq!(snapshot.terminal_without_report, vec![lost]); assert_eq!(snapshot.reported.len(), 1); } } diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index b41321ce3..746052a6f 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -6,6 +6,6 @@ mod prepare_static_plan; mod query_coordinator; mod store; -pub use distributed::{DistributedExec, TaskMetricsSnapshot}; +pub use distributed::DistributedExec; pub(crate) use dynamic_filter_registry::DynamicFilterRegistry; pub(crate) use store::Store; diff --git a/src/lib.rs b/src/lib.rs index f345502f4..523252982 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,7 +19,7 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::{DistributedExec, TaskMetricsSnapshot}; +pub use coordinator::DistributedExec; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, ProducerHead, SessionStateBuilderExt, diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 77ae75d79..e1ca7c088 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -179,7 +179,7 @@ pub struct TaskMetrics { impl TaskMetrics { /// Whether this report came from a planned task that never received ExecuteTask. Its plan /// may still contain real sampling metrics, so it must not be treated as a zero-cost task. - pub fn was_not_executed(&self) -> bool { + pub(crate) fn was_not_executed(&self) -> bool { self.task_metrics .iter() .any(|metric| metric.value().name() == crate::metrics::TASK_NOT_EXECUTED_METRIC) diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index 724e5647e..cdd003251 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -416,31 +416,10 @@ mod tests { ), ) .await??; - rewrite_with_metrics_within_timeout(with_filters).await?; - let snapshot = plan - .downcast_ref::() - .unwrap() - .metrics_snapshot() - .unwrap(); - assert!( - snapshot.all_reported(), - "unexpected lost report: {:?}", - snapshot.missing_reports - ); - let unexecuted = snapshot.unexecuted_tasks(); - assert!(!unexecuted.is_empty(), "AQE did not skip any task"); - assert!( - unexecuted.iter().any(|key| { - snapshot.reported[key] - .pre_order_plan_metrics - .iter() - .flat_map(|set| set.iter()) - .any(|metric| { - metric.value().name() == "files_opened" && metric.value().as_usize() > 0 - }) - }), - "skipped AQE tasks lost their actual sampling I/O metrics" - ); + let rewritten = rewrite_with_metrics_within_timeout(with_filters).await?; + let display = display_plan_ascii(rewritten.as_ref(), true); + assert_contains!(&display, "task_not_executed="); + assert_contains!(&display, "files_opened="); Ok(()) }