Metrics finalization can hang on planned but never executed tasks with AQE or static planning

[sesteves]

Summary

A query can finish execution successfully but hang indefinitely in
rewrite_distributed_plan_with_metrics when a planned worker task is never executed.

The original reliable reproducer used adaptive query execution (AQE), but production evidence
shows that AQE is not required. Static planning can trigger the same bug when execution
short-circuits, particularly with LIMIT.

Metrics collection must never prevent an otherwise successful query from returning.

Environment

• datafusion-distributed: 3.0.0
• DataFusion: 54.1.0
• collect_metrics = true
• Multi-worker execution
• Observed with both dynamic_task_count = true and false

The relevant behavior also appears to remain on current main.

Production Regression

Production impact increased sharply after upgrading from datafusion-distributed 1.0.0 at
revision b4856efac4b9bab696bfdac2c3dbe10b98934913 to 3.0.0.

However, code inspection shows that the previous pinned revision already contained the
underlying indefinite metrics wait and could omit metrics for a planned but never executed task.
Therefore, we cannot currently claim that 3.0.0 introduced the protocol bug itself.

The confirmed regression is that 3.0.0 makes this condition occur much more frequently, likely
because of changes to distributed planning and task execution. The exact behavioral change that
exposed the pre-existing metrics finalization bug has not yet been bisected.

Trigger 1: AQE

The original reliable reproducer uses a distributed hash join where:

• AQE installs and samples a remote producer stage.
• The hash build side is empty.
• DataFusion completes the join without polling the probe.
• A producer task receives SetPlanRequest but never receives ExecuteTask.
• Query result collection reaches EOS successfully.
• Metrics rewriting waits indefinitely for that task's missing metrics.

In the instrumented reproduction:

• Query execution completed in 33 ms.
• The request remained blocked until an outer timeout fired.
• Disabling metrics collection while keeping AQE enabled made the warmed reproducer pass 20/20
  times, compared with 6/12 hangs when metrics collection was enabled.
• Single-worker execution succeeded.

Trigger 2: Static Planning

Production requests continue to exhibit the same behavior with
dynamic_task_count = false.

A representative request had:

• An explicit optim_adaptive_query_execution=false override.
• A static distributed scan with LIMIT 1.
• 50 planned worker tasks.
• Planning completed in approximately 2.3 seconds.
• Result collection completed approximately 150 milliseconds later.
• Only enough scan work was performed to satisfy the limit.
• The request then remained blocked in metrics finalization until its outer timeout fired.

The relevant plan shape was:

CoalescePartitionsExec(fetch=1)
  ProjectionExec
    DataSourceExec(fetch=1)

Disabling AQE did not eliminate these production timeouts. More than 95% of affected production
queries had a LIMIT, consistent with execution completing before every planned task receives
ExecuteTask.

Root Cause

TaskData::final_plan is initialized only when ExecuteTask calls TaskData::plan:

https://github.com/datafusion-contrib/datafusion-distributed/blob/v3.0.0/src/worker/task_data.rs#L105-L123

When the coordinator channel reaches EOS, metrics are sent only if final_plan exists:

https://github.com/datafusion-contrib/datafusion-distributed/blob/v3.0.0/src/worker/impl_coordinator_channel.rs#L175-L187

A planned but never executed task therefore produces no TaskMetrics.

The coordinator derives expected task keys from the prepared plan and waits without a timeout for
every key:

https://github.com/datafusion-contrib/datafusion-distributed/blob/v3.0.0/src/coordinator/distributed.rs#L83-L110

Consequently, rewrite_distributed_plan_with_metrics waits forever:

https://github.com/datafusion-contrib/datafusion-distributed/blob/v3.0.0/src/metrics/task_metrics_rewriter.rs#L40-L53

AQE is one way to create a planned but never executed task. Static execution short-circuiting is
another.

Relationship To Existing Fixes

PR #524 moved metrics finalization to coordinator-channel EOS, but the EOS path still emits
nothing when final_plan was never initialized.

PR #482 and the alternative cancellation fix address abandoned network streams. They do not
guarantee that every planned task reaches a terminal metrics state.

Expected Behavior

Every planned task must reach a terminal metrics state, whether it was executed, failed,
cancelled, skipped, or never executed. Metrics collection must not block query completion
indefinitely.

Possible Fixes

• When the worker-to-coordinator stream ends without a TaskMetrics message, record an empty
  metric set for that task so it reaches a terminal state. This follows the approach suggested in
  Metrics finalization can hang on planned but never executed tasks with AQE or static planning #739 (comment).
• Represent skipped or never-executed tasks explicitly in the metrics store.
• Wait for every task to become terminal rather than requiring a non-empty metrics report from
  every planned task.
• Add a defensive timeout to metrics waiting so observability cannot deadlock query completion.

Consumers can use tokio::time::timeout around rewrite_distributed_plan_with_metrics as a
temporary workaround, but this can produce incomplete accounting and execution metrics.

Suggested Regression Tests

AQE

1. Enable adaptive task count and metrics collection.
2. Plan and sample a worker task.
3. Do not issue ExecuteTask for that task.
4. Close the coordinator channel.
5. Assert that metrics rewriting completes and represents the task as unexecuted or with empty
   metrics.

Static Planning

1. Disable adaptive task count and enable metrics collection.
2. Create a multi-worker distributed scan with substantially more tasks than needed to satisfy
   LIMIT 1.
3. Allow execution to short-circuit before every planned task is executed.
4. Assert that result collection and metrics rewriting both complete.
5. Assert that skipped tasks are terminal without fabricated execution metrics.

[gabotechs]

🤔 Interesting. I imagine this is also technically possible in static planning, but less likely to happen as there's not a significant time difference between the ExecuteTask calls in the different stages.

As a workaround in the mean time, you can wrap your rewrite_distributed_plan_with_metrics in a tokio::time::timeout in order to avoid deadlocks.

Thanks for the report!

[sesteves]

@gabotechs Thanks, that makes sense. We confirmed the hang is in metrics finalization: disabling metrics collection made the AQE reproducer pass 20/20 versus 6/12 hangs.

We’ll add a conservative timeout, though timed-out worker metrics may undercount usage telemetry. For the permanent fix, should DFD emit base_plan metrics when final_plan is unset, or represent the task as terminal-but-never-executed? Is base_plan compatible with plan_for_viz here?

[gabotechs]

Is base_plan compatible with plan_for_viz here?

Not really, base_plan is not going to give any visibility around what really happened during the query.

One thing that comes to mind would be to ensure that the spawned tokio task that is responsible for collecting metrics stores an empty metricset in case no TaskMetrics where received over the wire here:

https://github.com/datafusion-contrib/datafusion-distributed/blob/main/src/coordinator/query_coordinator.rs#L324-L350

If that while loop ends, and no TaskMetrics message was received, then it means the task was dropped and metrics will go missing.