Propagate span when spawning tasks

Used by CubeStore tracing infrastructure

Author: ilya-biryukov

rust/datafusion/src/cube_ext/mod.rs

@@ -21,3 +21,6 @@ pub mod joinagg;
 pub mod sequence;
 pub mod stream;
 pub mod util;
+
+mod spawn;
+pub use spawn::*;

rust/datafusion/src/cube_ext/spawn.rs

@@ -0,0 +1,48 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use futures::Future;
+use tokio::task::JoinHandle;
+use tracing_futures::Instrument;
+
+/// Calls [tokio::spawn] and additionally enables tracing of the spawned task as part of the current
+/// computation. This is CubeStore approach to tracing, so all code must use this function instead
+/// of replace [tokio::spawn].
+pub fn spawn<T>(task: T) -> JoinHandle<T::Output>
+where
+    T: Future + Send + 'static,
+    T::Output: Send + 'static,
+{
+    tokio::spawn(task.in_current_span())
+}
+
+/// Propagates current span to blocking operation. See [spawn] for details.
+pub fn spawn_blocking<F, R>(f: F) -> JoinHandle<R>
+where
+    F: FnOnce() -> R + Send + 'static,
+    R: Send + 'static,
+{
+    let span = tracing::Span::current();
+    if span.is_disabled() {
+        tokio::task::spawn_blocking(f)
+    } else {
+        tokio::task::spawn_blocking(move || {
+            let _ = tracing::info_span!(parent: &span, "blocking task").enter();
+            f()
+        })
+    }
+}

rust/datafusion/src/datasource/memory.rs

@@ -34,6 +34,7 @@ use crate::physical_plan::common;
 use crate::physical_plan::memory::MemoryExec;
 use crate::physical_plan::ExecutionPlan;
 use crate::{
+    cube_ext,
     datasource::datasource::Statistics,
     physical_plan::{repartition::RepartitionExec, Partitioning},
 };
@@ -121,7 +122,7 @@ impl MemTable {
         let tasks = (0..partition_count)
             .map(|part_i| {
                 let exec = exec.clone();
-                tokio::spawn(async move {
+                cube_ext::spawn(async move {
                     let stream = exec.execute(part_i).await?;
                     common::collect(stream).await
                 })

rust/datafusion/src/physical_plan/hash_aggregate.rs

@@ -77,6 +77,7 @@ use super::{
     SendableRecordBatchStream,
 };
 
+use crate::cube_ext;
 use crate::logical_plan::{DFSchema, DFSchemaRef};
 use crate::physical_plan::sorted_aggregate::SortedAggState;
 use crate::scalar::ScalarValue;
@@ -85,7 +86,6 @@ use itertools::Itertools;
 use smallvec::smallvec;
 use smallvec::SmallVec;
 use std::convert::TryFrom;
-use tracing_futures::{Instrument, WithSubscriber};
 
 /// Hash aggregate modes
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -765,7 +765,7 @@ impl GroupedHashAggregateStream {
             };
             tx.send(result)
         };
-        tokio::spawn(task.in_current_span().with_current_subscriber());
+        cube_ext::spawn(task);
 
         GroupedHashAggregateStream {
             schema,
@@ -945,7 +945,7 @@ impl HashAggregateStream {
                 compute_hash_aggregate(mode, schema_clone, aggr_expr, input).await;
             tx.send(result)
         };
-        tokio::spawn(task.in_current_span().with_current_subscriber());
+        cube_ext::spawn(task);
 
         HashAggregateStream {
             schema,

rust/datafusion/src/physical_plan/merge.rs

@@ -40,10 +40,10 @@ use crate::physical_plan::Partitioning;
 use crate::physical_plan::{ExecutionPlan, OptimizerHints};
 
 use super::SendableRecordBatchStream;
+use crate::cube_ext;
 use crate::logical_plan::DFSchemaRef;
 use pin_project_lite::pin_project;
 use std::option::Option::None;
-use tracing_futures::{Instrument, WithSubscriber};
 
 /// Merge execution plan executes partitions in parallel and combines them into a single
 /// partition. No guarantees are made about the order of the resulting partition.
@@ -145,7 +145,7 @@ impl ExecutionPlan for MergeExec {
                             sender.send(item).await.ok();
                         }
                     };
-                    tokio::spawn(task.in_current_span().with_current_subscriber());
+                    cube_ext::spawn(task);
                 }
 
                 Ok(Box::pin(MergeStream {

rust/datafusion/src/physical_plan/parquet.rs

@@ -32,6 +32,7 @@ use super::{
 };
 use crate::{
     catalog::catalog::MemoryCatalogList,
+    cube_ext,
     physical_plan::{common, ExecutionPlan, Partitioning},
 };
 use crate::{
@@ -59,10 +60,7 @@ use parquet::file::{
 
 use fmt::Debug;
 use parquet::arrow::{ArrowReader, ParquetFileArrowReader};
-use tokio::{
-    sync::mpsc::{channel, Receiver, Sender},
-    task,
-};
+use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio_stream::wrappers::ReceiverStream;
 
 use crate::datasource::datasource::{ColumnStatistics, Statistics};
@@ -877,7 +875,7 @@ impl ExecutionPlan for ParquetExec {
         let batch_size = self.batch_size;
         let limit = self.limit;
 
-        task::spawn_blocking(move || {
+        cube_ext::spawn_blocking(move || {
             if let Err(e) = read_files(
                 &filenames,
                 &projection,

rust/datafusion/src/physical_plan/repartition.rs

@@ -33,6 +33,7 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 use super::{hash_join::create_hashes, RecordBatchStream, SendableRecordBatchStream};
 use async_trait::async_trait;
 
+use crate::cube_ext;
 use crate::logical_plan::DFSchemaRef;
 use futures::stream::Stream;
 use futures::StreamExt;
@@ -142,7 +143,7 @@ impl ExecutionPlan for RepartitionExec {
                     .map(|(partition, (tx, _rx))| (*partition, tx.clone()))
                     .collect();
                 let partitioning = self.partitioning.clone();
-                let _: JoinHandle<Result<()>> = tokio::spawn(async move {
+                let _: JoinHandle<Result<()>> = cube_ext::spawn(async move {
                     let mut stream = input.execute(i).await?;
                     let mut counter = 0;
                     let hashes_buf = &mut vec![];
@@ -435,7 +436,7 @@ mod tests {
     #[tokio::test]
     async fn many_to_many_round_robin_within_tokio_task() -> Result<()> {
         let join_handle: JoinHandle<Result<Vec<Vec<RecordBatch>>>> =
-            tokio::spawn(async move {
+            cube_ext::spawn(async move {
                 // define input partitions
                 let schema = test_schema();
                 let partition = create_vec_batches(&schema, 50);

rust/datafusion/src/physical_plan/sort.rs

@@ -38,6 +38,7 @@ use arrow::record_batch::RecordBatch;
 use arrow::{array::ArrayRef, error::ArrowError};
 
 use super::{RecordBatchStream, SendableRecordBatchStream};
+use crate::cube_ext;
 use crate::error::{DataFusionError, Result};
 use crate::logical_plan::DFSchemaRef;
 use crate::physical_plan::expressions::PhysicalSortExpr;
@@ -249,7 +250,7 @@ impl SortStream {
         let (tx, rx) = futures::channel::oneshot::channel();
 
         let schema = input.schema();
-        tokio::spawn(async move {
+        cube_ext::spawn(async move {
             let schema = input.schema();
             let sorted_batch = common::collect(input)
                 .await