From 6f19f1d47db5fe9ec581199c8ae36e3d3c2e2e6d Mon Sep 17 00:00:00 2001
From: Luiz Carvalho <luizcmpc@gmail.com>
Date: Wed, 25 Sep 2024 14:05:14 -0300
Subject: [PATCH] fix(sentry-tracing): switch sentry spans on enter and exit

---
 sentry-core/src/performance.rs |  6 +++++
 sentry-tracing/src/layer.rs    | 43 ++++++++++++++++++++++++++++------
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/sentry-core/src/performance.rs b/sentry-core/src/performance.rs
index 062899ae..9209edca 100644
--- a/sentry-core/src/performance.rs
+++ b/sentry-core/src/performance.rs
@@ -667,6 +667,12 @@ impl Span {
         transaction.context.clone()
     }
 
+    /// Get the current span ID.
+    pub fn get_span_id(&self) -> protocol::SpanId {
+        let span = self.span.lock().unwrap();
+        span.span_id
+    }
+
     /// Get the status of the Span.
     pub fn get_status(&self) -> Option<protocol::SpanStatus> {
         let span = self.span.lock().unwrap();
diff --git a/sentry-tracing/src/layer.rs b/sentry-tracing/src/layer.rs
index 8183bbfc..6790e422 100644
--- a/sentry-tracing/src/layer.rs
+++ b/sentry-tracing/src/layer.rs
@@ -221,8 +221,6 @@ where
             sentry_span.set_data(key, value);
         }
 
-        sentry_core::configure_scope(|scope| scope.set_span(Some(sentry_span.clone())));
-
         let mut extensions = span.extensions_mut();
         extensions.insert(SentrySpanData {
             sentry_span,
@@ -230,6 +228,41 @@ where
         });
     }
 
+    /// Sets entered span as *current* sentry span. A tracing span can be
+    /// entered and existed multiple times, for example, when using a `tracing::Instrumented` future.
+    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
+        let span = match ctx.span(&id) {
+            Some(span) => span,
+            None => return,
+        };
+
+        let extensions = span.extensions();
+        let SentrySpanData { sentry_span, .. } = match extensions.get::<SentrySpanData>() {
+            Some(data) => data,
+            None => return,
+        };
+
+        sentry_core::configure_scope(|scope| scope.set_span(Some(sentry_span.clone())));
+    }
+
+    /// Set exited span's parent as *current* sentry span.
+    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
+        let span = match ctx.span(&id) {
+            Some(span) => span,
+            None => return,
+        };
+
+        let extensions = span.extensions();
+        let SentrySpanData {
+            parent_sentry_span, ..
+        } = match extensions.get::<SentrySpanData>() {
+            Some(data) => data,
+            None => return,
+        };
+
+        sentry_core::configure_scope(|scope| scope.set_span(parent_sentry_span.clone()));
+    }
+
     /// When a span gets closed, finish the underlying sentry span, and set back
     /// its parent as the *current* sentry span.
     fn on_close(&self, id: span::Id, ctx: Context<'_, S>) {
@@ -239,16 +272,12 @@ where
         };
 
         let mut extensions = span.extensions_mut();
-        let SentrySpanData {
-            sentry_span,
-            parent_sentry_span,
-        } = match extensions.remove::<SentrySpanData>() {
+        let SentrySpanData { sentry_span, .. } = match extensions.remove::<SentrySpanData>() {
             Some(data) => data,
             None => return,
         };
 
         sentry_span.finish();
-        sentry_core::configure_scope(|scope| scope.set_span(parent_sentry_span));
     }
 
     /// Implement the writing of extra data to span