From f537a943907ce7dd375b3ebb56a6bdb3553385a1 Mon Sep 17 00:00:00 2001
From: DoumanAsh <douman@gmx.se>
Date: Thu, 31 Oct 2024 12:25:51 +0900
Subject: [PATCH] Add hub switch guard to the tracing layer

---
 sentry-core/src/hub.rs      |  7 ++++++
 sentry-core/src/hub_impl.rs |  5 ++++-
 sentry-core/src/lib.rs      |  1 +
 sentry-tracing/src/layer.rs | 43 +++++++++++++++++++++----------------
 4 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/sentry-core/src/hub.rs b/sentry-core/src/hub.rs
index f5ebf155..456948a0 100644
--- a/sentry-core/src/hub.rs
+++ b/sentry-core/src/hub.rs
@@ -4,6 +4,7 @@
 
 use std::sync::{Arc, RwLock};
 
+use crate::hub_impl::SwitchGuard;
 use crate::protocol::{Event, Level, SessionStatus};
 use crate::types::Uuid;
 use crate::{Integration, IntoBreadcrumbs, Scope, ScopeGuard};
@@ -242,4 +243,10 @@ impl Hub {
             })
         }}
     }
+
+    #[inline(always)]
+    /// Consumes self to create switch guard, that can be used to modify current hub in thread local storage
+    pub fn into_switch_guard(self: Arc<Hub>) -> SwitchGuard {
+        crate::hub_impl::SwitchGuard::new(self)
+    }
 }
diff --git a/sentry-core/src/hub_impl.rs b/sentry-core/src/hub_impl.rs
index 1a2de2cd..6692d1a1 100644
--- a/sentry-core/src/hub_impl.rs
+++ b/sentry-core/src/hub_impl.rs
@@ -21,7 +21,10 @@ thread_local! {
     );
 }
 
-pub(crate) struct SwitchGuard {
+///Hub switch guard
+///
+///Used to temporarily swap active hub in thread local storage.
+pub struct SwitchGuard {
     inner: Option<(Arc<Hub>, bool)>,
 }
 
diff --git a/sentry-core/src/lib.rs b/sentry-core/src/lib.rs
index 1f7d8c31..b72f25f6 100644
--- a/sentry-core/src/lib.rs
+++ b/sentry-core/src/lib.rs
@@ -130,6 +130,7 @@ pub use crate::clientoptions::{BeforeCallback, ClientOptions, SessionMode};
 pub use crate::error::{capture_error, event_from_error, parse_type_from_debug};
 pub use crate::futures::{SentryFuture, SentryFutureExt};
 pub use crate::hub::Hub;
+pub use crate::hub_impl::SwitchGuard as HubSwitchGuard;
 pub use crate::integration::Integration;
 pub use crate::intodsn::IntoDsn;
 pub use crate::performance::*;
diff --git a/sentry-tracing/src/layer.rs b/sentry-tracing/src/layer.rs
index b6acab31..e48fd71d 100644
--- a/sentry-tracing/src/layer.rs
+++ b/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 => {
@@ -272,42 +277,42 @@ where
         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) {
+        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())));
+        let mut extensions = span.extensions_mut();
+        if let Some(data) = extensions.get_mut::<SentrySpanData>() {
+            data.hub_switch_guard = Some(data.hub.clone().into_switch_guard());
+            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) {
+        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()));
+        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 = None;
+        }
     }
 
     /// When a span gets closed, finish the underlying sentry span, and set back