diff --git a/sentry-core/src/hub_impl.rs b/sentry-core/src/hub_impl.rs index 1a2de2cd..59c91bf9 100644 --- a/sentry-core/src/hub_impl.rs +++ b/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, bool)>, } impl SwitchGuard { - pub(crate) fn new(mut hub: Arc) -> 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) -> 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. diff --git a/sentry-core/src/lib.rs b/sentry-core/src/lib.rs index 1f7d8c31..5cc39cd5 100644 --- a/sentry-core/src/lib.rs +++ b/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")] diff --git a/sentry-core/src/performance.rs b/sentry-core/src/performance.rs index 59cca2a5..140d3120 100644 --- a/sentry-core/src/performance.rs +++ b/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 { let span = self.span.lock().unwrap(); diff --git a/sentry-tracing/src/layer.rs b/sentry-tracing/src/layer.rs index 01930f80..0540fe76 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 + Into>>( pub(super) struct SentrySpanData { pub(super) sentry_span: TransactionOrSpan, parent_sentry_span: Option, + hub: Arc, + hub_switch_guard: Option, } impl Layer for SentryLayer @@ -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::() { + 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::() { + 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::() { + let SentrySpanData { sentry_span, .. } = match extensions.remove::() { 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