fix(sentry-tracing): switch sentry spans on enter and exit (#686)

---------

Co-authored-by: DoumanAsh <[email protected]>

Author: saiintbrisson

sentry-core/src/hub_impl.rs

@@ -21,12 +21,17 @@ thread_local! {
     );
 }
 
-pub(crate) struct SwitchGuard {
+/// A Hub switch guard used to temporarily swap
+/// active hub in thread local storage.
+pub struct SwitchGuard {
     inner: Option<(Arc<Hub>, bool)>,
 }
 
 impl SwitchGuard {
-    pub(crate) fn new(mut hub: Arc<Hub>) -> Self {
+    /// Swaps the current thread's Hub by the one provided
+    /// and returns a guard that, when dropped, replaces it
+    /// to the previous one.
+    pub fn new(mut hub: Arc<Hub>) -> Self {
         let inner = THREAD_HUB.with(|(thread_hub, is_process_hub)| {
             // SAFETY: `thread_hub` will always be a valid thread local hub,
             // by definition not shared between threads.

sentry-core/src/lib.rs

@@ -149,8 +149,9 @@ pub mod metrics;
 mod session;
 #[cfg(all(feature = "client", feature = "metrics"))]
 mod units;
+
 #[cfg(feature = "client")]
-pub use crate::client::Client;
+pub use crate::{client::Client, hub_impl::SwitchGuard as HubSwitchGuard};
 
 // test utilities
 #[cfg(feature = "test")]

sentry-core/src/performance.rs

@@ -717,6 +717,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();

sentry-tracing/src/layer.rs

@@ -1,6 +1,7 @@
 use std::borrow::Cow;
 use std::cell::RefCell;
 use std::collections::BTreeMap;
+use std::sync::Arc;
 
 use sentry_core::protocol::Value;
 use sentry_core::{Breadcrumb, TransactionOrSpan};
@@ -197,6 +198,8 @@ fn record_fields<'a, K: AsRef<str> + Into<Cow<'a, str>>>(
 pub(super) struct SentrySpanData {
     pub(super) sentry_span: TransactionOrSpan,
     parent_sentry_span: Option<TransactionOrSpan>,
+    hub: Arc<sentry_core::Hub>,
+    hub_switch_guard: Option<sentry_core::HubSwitchGuard>,
 }
 
 impl<S> Layer<S> for SentryLayer<S>
@@ -256,7 +259,9 @@ where
             }
         });
 
-        let parent_sentry_span = sentry_core::configure_scope(|s| s.get_span());
+        let hub = sentry_core::Hub::current();
+        let parent_sentry_span = hub.configure_scope(|scope| scope.get_span());
+
         let sentry_span: sentry_core::TransactionOrSpan = match &parent_sentry_span {
             Some(parent) => parent.start_child(op, &description).into(),
             None => {
@@ -268,15 +273,48 @@ where
         // This comes from typically the `fields` in `tracing::instrument`.
         record_fields(&sentry_span, data);
 
-        sentry_core::configure_scope(|scope| scope.set_span(Some(sentry_span.clone())));
-
         let mut extensions = span.extensions_mut();
         extensions.insert(SentrySpanData {
             sentry_span,
             parent_sentry_span,
+            hub,
+            hub_switch_guard: None,
         });
     }
 
+    /// 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 mut extensions = span.extensions_mut();
+        if let Some(data) = extensions.get_mut::<SentrySpanData>() {
+            data.hub_switch_guard = Some(sentry_core::HubSwitchGuard::new(data.hub.clone()));
+            data.hub.configure_scope(|scope| {
+                scope.set_span(Some(data.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 mut extensions = span.extensions_mut();
+        if let Some(data) = extensions.get_mut::<SentrySpanData>() {
+            data.hub.configure_scope(|scope| {
+                scope.set_span(data.parent_sentry_span.clone());
+            });
+            data.hub_switch_guard.take();
+        }
+    }
+
     /// 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>) {
@@ -286,16 +324,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