add tokio_metrics

coralogix · Sep 21, 2023 · ace60bb · ace60bb
1 parent 63adfb6
commit ace60bb
Show file tree

Hide file tree

Showing 9 changed files with 14,949 additions and 109 deletions.
diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml
@@ -95,6 +95,7 @@ uuid = { version = "1.0", features = ["v4"] }
 xz2 = { version = "0.1", optional = true }
 zstd = { version = "0.12", optional = true, default-features = false }
 
+tokio-metrics = "0.3.0"
 
 [dev-dependencies]
 async-trait = "0.1.53"

diff --git a/datafusion/core/src/physical_plan/common.rs b/datafusion/core/src/physical_plan/common.rs
@@ -98,6 +98,7 @@ fn build_file_list_recurse(
 pub(crate) fn spawn_buffered(
     mut input: SendableRecordBatchStream,
     buffer: usize,
+    monitor: Option<Arc<tokio_metrics::TaskMonitor>>,
 ) -> SendableRecordBatchStream {
     // Use tokio only if running from a tokio context (#2201)
     if tokio::runtime::Handle::try_current().is_err() {
@@ -108,13 +109,19 @@ pub(crate) fn spawn_buffered(
 
     let sender = builder.tx();
 
-    builder.spawn(async move {
+    let task = async move {
         while let Some(item) = input.next().await {
             if sender.send(item).await.is_err() {
                 return;
             }
         }
-    });
+    };
+
+    if let Some(monitor) = monitor {
+        builder.spawn(monitor.instrument(task));
+    } else {
+        builder.spawn(task);
+    }
 
     builder.build()
 }

diff --git a/datafusion/core/src/physical_plan/repartition/mod.rs b/datafusion/core/src/physical_plan/repartition/mod.rs
@@ -473,6 +473,10 @@ impl ExecutionPlan for RepartitionExec {
             self.name(),
             partition
         );
+
+        let monitor = context
+            .session_config()
+            .get_extension::<tokio_metrics::TaskMonitor>();
         // lock mutexes
         let mut state = self.state.lock();
 
@@ -522,15 +526,20 @@ impl ExecutionPlan for RepartitionExec {
 
                 let r_metrics = RepartitionMetrics::new(i, partition, &self.metrics);
 
+                let task = Self::pull_from_input(
+                    self.input.clone(),
+                    i,
+                    txs.clone(),
+                    self.partitioning.clone(),
+                    r_metrics,
+                    context.clone(),
+                );
                 let input_task: JoinHandle<Result<()>> =
-                    tokio::spawn(Self::pull_from_input(
-                        self.input.clone(),
-                        i,
-                        txs.clone(),
-                        self.partitioning.clone(),
-                        r_metrics,
-                        context.clone(),
-                    ));
+                    if let Some(monitor) = monitor.as_deref() {
+                        tokio::spawn(monitor.instrument(task))
+                    } else {
+                        tokio::spawn(task)
+                    };
 
                 // In a separate task, wait for each input to be done
                 // (and pass along any errors, including panic!s)

diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs b/datafusion/core/src/physical_plan/sorts/sort.rs
@@ -282,7 +282,11 @@ impl ExternalSorter {
             .into_iter()
             .map(|batch| {
                 let metrics = self.metrics.baseline.intermediate();
-                Ok(spawn_buffered(self.sort_batch_stream(batch, metrics)?, 1))
+                Ok(spawn_buffered(
+                    self.sort_batch_stream(batch, metrics)?,
+                    1,
+                    None,
+                ))
             })
             .collect::<Result<_>>()?;
 

diff --git a/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs b/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs
@@ -39,6 +39,7 @@ use datafusion_physical_expr::{
 };
 
 use log::{debug, trace};
+use tokio_metrics::TaskMonitor;
 
 /// Sort preserving merge execution plan
 ///
@@ -193,6 +194,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
             "Start SortPreservingMergeExec::execute for partition: {}",
             partition
         );
+        let monitor = context.session_config().get_extension::<TaskMonitor>();
         if 0 != partition {
             return Err(DataFusionError::Internal(format!(
                 "SortPreservingMergeExec invalid partition {partition}"
@@ -221,7 +223,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
                 let receivers = (0..input_partitions)
                     .map(|partition| {
                         let stream = self.input.execute(partition, context.clone())?;
-                        Ok(spawn_buffered(stream, 1))
+                        Ok(spawn_buffered(stream, 1, monitor.clone()))
                     })
                     .collect::<Result<_>>()?;
 

diff --git a/datafusion/core/src/physical_plan/stream.rs b/datafusion/core/src/physical_plan/stream.rs
@@ -106,7 +106,11 @@ impl RecordBatchReceiverStreamBuilder {
     ) {
         let output = self.tx();
 
-        self.spawn(async move {
+        let monitor = context
+            .session_config()
+            .get_extension::<tokio_metrics::TaskMonitor>();
+
+        let task = async move {
             let mut stream = match input.execute(partition, context) {
                 Err(e) => {
                     // If send fails, the plan being torn down, there
@@ -146,7 +150,13 @@ impl RecordBatchReceiverStreamBuilder {
                     return;
                 }
             }
-        });
+        };
+
+        if let Some(monitor) = monitor {
+            self.spawn(monitor.instrument(task));
+        } else {
+            self.spawn(task);
+        }
     }
 
     /// Create a stream of all `RecordBatch`es written to `tx`