feat(cube): Pass tracing spans through spawned tasks in ExecutionPlan execution

Author: srh

datafusion/physical-plan/src/cube_ext/spawn.rs

@@ -19,8 +19,9 @@ use std::future::Future;
 use crate::cube_ext::catch_unwind::{
     async_try_with_catch_unwind, try_with_catch_unwind, PanicError,
 };
+use datafusion_common_runtime::SpawnedTask;
 use futures::sink::SinkExt;
-use tokio::task::JoinHandle;
+use tokio::task::{AbortHandle, JoinHandle, JoinSet};
 use tracing_futures::Instrument;
 
 /// Calls [tokio::spawn] and additionally enables tracing of the spawned task as part of the current
@@ -57,6 +58,56 @@ where
     }
 }
 
+/// Propagates current span to the spawned task.  See [spawn] for details.
+pub fn spawn_on_joinset<F, T>(join_set: &mut JoinSet<T>, task: F) -> AbortHandle
+    where
+        F: Future<Output = T>,
+        F: Send + 'static,
+        T: Send + 'static,
+{
+    if let Some(s) = new_subtask_span() {
+        join_set.spawn(async move {
+            let _p = s.parent;  // ensure parent stays alive.
+            task.instrument(s.child).await
+        })
+    } else {
+        join_set.spawn(task)
+    }
+}
+
+/// Propagates current span to the blocking operation.  See [spawn] for details.
+pub fn spawn_blocking_on_joinset<F, T>(join_set: &mut JoinSet<T>, f: F) -> AbortHandle
+    where
+        F: FnOnce() -> T,
+        F: Send + 'static,
+        T: Send + 'static,
+{
+    if let Some(s) = new_subtask_span() {
+        join_set.spawn_blocking(move || {
+            let _p = s.parent;  // ensure parent stays alive.
+            s.child.in_scope(f)
+        })
+    } else {
+        join_set.spawn_blocking(f)
+    }
+}
+
+/// Propagates current span to the spawned task.  See [spawn] for details.
+pub fn spawn_spawned_task<T>(task: T) -> SpawnedTask<T::Output>
+where
+    T: Future + Send + 'static,
+    T::Output: Send + 'static,
+{
+    if let Some(s) = new_subtask_span() {
+        SpawnedTask::spawn(async move {
+            let _p = s.parent; // ensure parent stays alive.
+            task.instrument(s.child).await
+        })
+    } else {
+        SpawnedTask::spawn(task)
+    }
+}
+
 struct SpawnSpans {
     parent: tracing::Span,
     child: tracing::Span,

datafusion/physical-plan/src/repartition/mod.rs

@@ -133,7 +133,7 @@ impl RepartitionExecState {
 
             let r_metrics = RepartitionMetrics::new(i, num_output_partitions, &metrics);
 
-            let input_task = SpawnedTask::spawn(RepartitionExec::pull_from_input(
+            let input_task = crate::cube_ext::spawn_spawned_task(RepartitionExec::pull_from_input(
                 Arc::clone(&input),
                 i,
                 txs.clone(),
@@ -144,7 +144,7 @@ impl RepartitionExecState {
 
             // In a separate task, wait for each input to be done
             // (and pass along any errors, including panic!s)
-            let wait_for_task = SpawnedTask::spawn(RepartitionExec::wait_for_task(
+            let wait_for_task = crate::cube_ext::spawn_spawned_task(RepartitionExec::wait_for_task(
                 input_task,
                 txs.into_iter()
                     .map(|(partition, (tx, _reservation))| (partition, tx))

datafusion/physical-plan/src/stream.rs

@@ -79,7 +79,7 @@ impl<O: Send + 'static> ReceiverStreamBuilder<O> {
         F: Future<Output = Result<()>>,
         F: Send + 'static,
     {
-        self.join_set.spawn(task);
+        crate::cube_ext::spawn_on_joinset(&mut self.join_set, task);
     }
 
     /// Spawn a blocking task that will be aborted if this builder (or the stream
@@ -92,7 +92,7 @@ impl<O: Send + 'static> ReceiverStreamBuilder<O> {
         F: FnOnce() -> Result<()>,
         F: Send + 'static,
     {
-        self.join_set.spawn_blocking(f);
+        crate::cube_ext::spawn_blocking_on_joinset(&mut self.join_set, f);
     }
 
     /// Create a stream of all data written to `tx`