fix: await async metric delivery in metrics collection tests (#490)

## Which issue does this PR close?

Closes #487.

## Rationale for this change

`test_metrics_collection_e2e_1` flakes intermittently in CI, panicking
on `Option::unwrap()` of a `None` value at the task-key lookup in
`run_metrics_collection_e2e_test` (e.g. [this
run](https://github.com/datafusion-contrib/datafusion-distributed/actions/runs/27137166880/job/80101529756)).

Per-task metrics are delivered to the coordinator asynchronously over
the `WorkerToCoordinator` side channel by a background task, which may
still be in flight after `execute_plan` returns. The test asserted on
the collected metrics immediately, racing that delivery, so a task key
was occasionally still absent from the `MetricsStore` when the assertion
ran.

## What changes are included in this PR?

- Add a `wait_for_task_metrics` test helper that waits on the
`MetricsStore` watch channel until every expected task key is present
before asserting, bounded by a 10s timeout that reports any
still-missing keys. This mirrors the existing
`DistributedExec::wait_for_metrics` pattern but waits on exactly the
task keys the assertions check.
- Use it in `run_metrics_collection_e2e_test` (covers `_e2e_1` and
`_e2e_2`).
- `test_metrics_collection_with_limit_causing_early_stream_termination`
had the same race papered over with a single `yield_now().await`;
replace that with the same deterministic wait.

## Are there any user-facing changes?

No. Test-only change.

---------

Signed-off-by: sjhddh <151469562+sjhddh@users.noreply.github.com>
Co-authored-by: sjhddh <151469562+sjhddh@users.noreply.github.com>

Author: sjhddh

src/metrics/task_metrics_collector.rs

@@ -172,14 +172,15 @@ mod tests {
             DisplayableExecutionPlan::new(plan.as_ref()).indent(true)
         );
 
+        // Per-task metrics are delivered asynchronously over the `WorkerToCoordinator` side
+        // channel after execution completes; await that delivery instead of racing it (see #487).
+        dist_exec.wait_for_metrics().await;
+
+        let metrics_store = dist_exec.metrics_store.as_ref().unwrap();
+
         // Ensure that there's metrics for each node for each task for each stage.
         for expected_task_key in expected_task_keys {
-            let actual_metrics = dist_exec
-                .metrics_store
-                .as_ref()
-                .unwrap()
-                .get(&expected_task_key)
-                .unwrap();
+            let actual_metrics = metrics_store.get(&expected_task_key).unwrap();
 
             // Verify that metrics were collected for all nodes. Some nodes may legitimately have
             // empty metrics (e.g., custom execution plans without metrics), which is fine - we
@@ -294,22 +295,18 @@ mod tests {
         execute_plan(plan.clone(), &ctx).await;
 
         // Metrics are delivered via the WorkerToCoordinator side channel in a background task.
-        // Yield briefly to let it complete before asserting.
-        tokio::task::yield_now().await;
+        // Wait for that delivery to complete before asserting, rather than racing it.
+        dist_exec.wait_for_metrics().await;
+        let metrics_store = dist_exec.metrics_store.as_ref().unwrap();
 
         for expected_task_key in &expected_task_keys {
-            let actual_metrics = dist_exec
-                .metrics_store
-                .as_ref()
-                .unwrap()
-                .get(expected_task_key)
-                .unwrap_or_else(|| {
-                    panic!(
-                        "Missing metrics for task key {expected_task_key:?}. \
+            let actual_metrics = metrics_store.get(expected_task_key).unwrap_or_else(|| {
+                panic!(
+                    "Missing metrics for task key {expected_task_key:?}. \
                          The LIMIT caused the stream to be dropped before the worker \
                          sent metrics via the coordinator channel."
-                    )
-                });
+                )
+            });
             let stage = stages.get(&(expected_task_key.stage_id as usize)).unwrap();
             let stage_plan = stage.local_plan().unwrap();
             assert_eq!(