feat: When tracing is enabled use the span from the scope of the run

aftershootco · Feb 28, 2025 · 560224b · 560224b
1 parent 4daae13
commit 560224b
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/mnn-sync/src/lib.rs b/mnn-sync/src/lib.rs
@@ -365,9 +365,13 @@ impl SessionHandle {
         f: impl FnOnce(&mut SessionRunner) -> Result<R> + Send + Sync + 'static,
     ) -> Result<R> {
         self.ensure_running()?;
+        #[cfg(features = "tracing")]
+        let span = tracing::Span::current();
         let f = f;
         let (tx, rx) = oneshot::channel();
         let wrapped_f = move |sr: &mut SessionRunner| -> Result<()> {
+            #[cfg(features = "tracing")]
+            let _guard = span.enter();
             let result = f(sr);
             tx.send(result)
                 .change_context(ErrorKind::SyncError)
@@ -387,9 +391,13 @@ impl SessionHandle {
         f: impl FnOnce(&mut SessionRunner) -> Result<R> + Send + Sync + 'static,
     ) -> Result<R> {
         self.ensure_running()?;
+        #[cfg(features = "tracing")]
+        let span = tracing::Span::current();
         let f = f;
         let (tx, rx) = oneshot::channel();
         let wrapped_f = move |sr: &mut SessionRunner| -> Result<()> {
+            #[cfg(features = "tracing")]
+            let _guard = span.enter();
             let result = f(sr);
             tx.send(result)
                 .change_context(ErrorKind::SyncError)