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

Author: saiintbrisson

Cargo.lock

@@ -322,6 +322,24 @@ impl TransactionOrSpan {
         }
     }
 
+    #[doc(hidden)]
+    #[must_use = "a span must be explicitly closed via `finish()`"]
+    pub fn start_child_with_span_id(
+        &self,
+        op: &str,
+        description: &str,
+        span_id: protocol::SpanId,
+    ) -> Span {
+        match self {
+            TransactionOrSpan::Transaction(transaction) => {
+                transaction.start_child_with_span_id(op, description, span_id)
+            }
+            TransactionOrSpan::Span(span) => {
+                span.start_child_with_span_id(op, description, span_id)
+            }
+        }
+    }
+
     #[cfg(feature = "client")]
     pub(crate) fn apply_to_event(&self, event: &mut protocol::Event<'_>) {
         if event.contexts.contains_key("trace") {
@@ -587,10 +605,22 @@ impl Transaction {
     /// The span must be explicitly finished via [`Span::finish`].
     #[must_use = "a span must be explicitly closed via `finish()`"]
     pub fn start_child(&self, op: &str, description: &str) -> Span {
+        self.start_child_with_span_id(op, description, Default::default())
+    }
+
+    #[doc(hidden)]
+    #[must_use = "a span must be explicitly closed via `finish()`"]
+    pub fn start_child_with_span_id(
+        &self,
+        op: &str,
+        description: &str,
+        span_id: protocol::SpanId,
+    ) -> Span {
         let inner = self.inner.lock().unwrap();
         let span = protocol::Span {
             trace_id: inner.context.trace_id,
             parent_span_id: Some(inner.context.span_id),
+            span_id,
             op: Some(op.into()),
             description: if description.is_empty() {
                 None
@@ -754,10 +784,22 @@ impl Span {
     /// The span must be explicitly finished via [`Span::finish`].
     #[must_use = "a span must be explicitly closed via `finish()`"]
     pub fn start_child(&self, op: &str, description: &str) -> Span {
+        self.start_child_with_span_id(op, description, Default::default())
+    }
+
+    #[doc(hidden)]
+    #[must_use = "a span must be explicitly closed via `finish()`"]
+    pub fn start_child_with_span_id(
+        &self,
+        op: &str,
+        description: &str,
+        span_id: protocol::SpanId,
+    ) -> Span {
         let span = self.span.lock().unwrap();
         let span = protocol::Span {
             trace_id: span.trace_id,
             parent_span_id: Some(span.span_id),
+            span_id,
             op: Some(op.into()),
             description: if description.is_empty() {
                 None

sentry-core/src/performance.rs

@@ -208,8 +208,11 @@ where
         });
 
         let parent_sentry_span = sentry_core::configure_scope(|s| s.get_span());
+        let span_id = sentry_core::protocol::SpanId::from(id.into_u64().to_ne_bytes());
         let sentry_span: sentry_core::TransactionOrSpan = match &parent_sentry_span {
-            Some(parent) => parent.start_child(op, &description).into(),
+            Some(parent) => parent
+                .start_child_with_span_id(op, &description, span_id)
+                .into(),
             None => {
                 let ctx = sentry_core::TransactionContext::new(&description, op);
                 sentry_core::start_transaction(ctx).into()
@@ -221,15 +224,48 @@ 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,
             parent_sentry_span,
         });
     }
 
+    /// 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 +275,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

sentry-tracing/src/layer.rs

@@ -1355,6 +1355,12 @@ impl From<SpanId> for String {
     }
 }
 
+impl From<[u8; 8]> for SpanId {
+    fn from(span_id: [u8; 8]) -> Self {
+        Self(span_id)
+    }
+}
+
 impl str::FromStr for SpanId {
     type Err = hex::FromHexError;