fix issue #946

Author: szokeasaurusrex

sentry-tracing/src/layer.rs sentry-tracing/src/layer/mod.rssentry-tracing/src/layer.rs renamed to sentry-tracing/src/layer/mod.rs

@@ -1,11 +1,11 @@
 use std::borrow::Cow;
 use std::cell::RefCell;
-use std::collections::{BTreeMap, HashMap};
+use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use bitflags::bitflags;
 use sentry_core::protocol::Value;
-use sentry_core::{Breadcrumb, HubSwitchGuard, TransactionOrSpan};
+use sentry_core::{Breadcrumb, Hub, HubSwitchGuard, TransactionOrSpan};
 use tracing_core::field::Visit;
 use tracing_core::{span, Event, Field, Level, Metadata, Subscriber};
 use tracing_subscriber::layer::{Context, Layer};
@@ -16,6 +16,9 @@ use crate::SENTRY_NAME_FIELD;
 use crate::SENTRY_OP_FIELD;
 use crate::SENTRY_TRACE_FIELD;
 use crate::TAGS_PREFIX;
+use span_guard_stack::SpanGuardStack;
+
+mod span_guard_stack;
 
 bitflags! {
     /// The action that Sentry should perform for a given [`Event`]
@@ -350,19 +353,20 @@ where
 
         let extensions = span.extensions();
         if let Some(data) = extensions.get::<SentrySpanData>() {
-            let guard = HubSwitchGuard::new(data.hub.clone());
+            let hub = Arc::new(Hub::new_from_top(data.hub.clone()));
+            let guard = HubSwitchGuard::new(hub.clone());
             SPAN_GUARDS.with(|guards| {
-                guards.borrow_mut().insert(id.clone(), guard);
+                guards.borrow_mut().push(id.clone(), guard);
             });
-            data.hub.configure_scope(|scope| {
+            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 _guard = SPAN_GUARDS.with(|guards| guards.borrow_mut().remove(id));
+        let _guard = SPAN_GUARDS.with(|guards| guards.borrow_mut().pop(id.clone()));
 
         let span = match ctx.span(id) {
             Some(span) => span,
@@ -511,8 +515,7 @@ thread_local! {
     static VISITOR_BUFFER: RefCell<String> = const { RefCell::new(String::new()) };
     /// Hub switch guards keyed by span ID. Stored in thread-local so guards are
     /// always dropped on the same thread where they were created.
-    static SPAN_GUARDS: RefCell<HashMap<span::Id, HubSwitchGuard>> =
-        RefCell::new(HashMap::new());
+    static SPAN_GUARDS: RefCell<SpanGuardStack> = RefCell::new(SpanGuardStack::new());
 }
 
 /// Records all span fields into a `BTreeMap`, reusing a mutable `String` as buffer.

sentry-tracing/src/layer/span_guard_stack.rs

@@ -0,0 +1,118 @@
+//! This module contains code for stack-like storage for `HubSwitchGuard`s keyed
+//! by tracing span ID.
+
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
+
+use sentry_core::HubSwitchGuard;
+use tracing_core::span::Id as SpanId;
+
+/// Holds per-span stacks of `HubSwitchGuard`s to handle span re-entrancy.
+///
+/// Each time a span is entered, we should push a new guard onto the stack.
+/// When the span exits, we should pop the guard from the stack.
+pub(super) struct SpanGuardStack {
+    /// The map of span IDs to their respective guard stacks.
+    guards: HashMap<SpanId, Vec<HubSwitchGuard>>,
+}
+
+impl SpanGuardStack {
+    /// Creates an empty guard stack map.
+    pub(super) fn new() -> Self {
+        Self {
+            guards: HashMap::new(),
+        }
+    }
+
+    /// Pushes a guard for the given span ID, creating the stack if needed.
+    pub(super) fn push(&mut self, id: SpanId, guard: HubSwitchGuard) {
+        self.guards.entry(id).or_default().push(guard);
+    }
+
+    /// Pops the most recent guard for the span ID, removing the stack when empty.
+    pub(super) fn pop(&mut self, id: SpanId) -> Option<HubSwitchGuard> {
+        match self.guards.entry(id) {
+            Entry::Occupied(mut entry) => {
+                let stack = entry.get_mut();
+                let guard = stack.pop();
+                if stack.is_empty() {
+                    entry.remove();
+                }
+                guard
+            }
+            Entry::Vacant(_) => None,
+        }
+    }
+
+    /// Removes all guards for the span ID without returning them.
+    ///
+    /// This function guarantees that the guards are dropped in LIFO order.
+    /// That way, the hub which was active when the span was first entered
+    /// will be the one active after this function returns.
+    ///
+    /// Typically, remove should only get called once the span is fully
+    /// exited, so this removal order guarantee is mostly just defensive.
+    pub(super) fn remove(&mut self, id: &SpanId) {
+        self.guards
+            .remove(id)
+            .into_iter()
+            .flatten()
+            .rev() // <- we drop in reverse order
+            .for_each(drop);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::SpanGuardStack;
+    use sentry_core::{Hub, HubSwitchGuard};
+    use std::sync::Arc;
+    use tracing_core::span::Id as SpanId;
+
+    #[test]
+    fn pop_is_lifo() {
+        let initial = Hub::current();
+        let hub_a = Arc::new(Hub::new_from_top(initial.clone()));
+        let hub_b = Arc::new(Hub::new_from_top(hub_a.clone()));
+
+        let mut stack = SpanGuardStack::new();
+        let id = SpanId::from_u64(1);
+
+        stack.push(id.clone(), HubSwitchGuard::new(hub_a.clone()));
+        assert!(Arc::ptr_eq(&Hub::current(), &hub_a));
+
+        stack.push(id.clone(), HubSwitchGuard::new(hub_b.clone()));
+        assert!(Arc::ptr_eq(&Hub::current(), &hub_b));
+
+        drop(stack.pop(id.clone()).expect("guard for hub_b"));
+        assert!(Arc::ptr_eq(&Hub::current(), &hub_a));
+
+        drop(stack.pop(id.clone()).expect("guard for hub_a"));
+        assert!(Arc::ptr_eq(&Hub::current(), &initial));
+
+        assert!(stack.pop(id).is_none());
+    }
+
+    #[test]
+    fn remove_drops_all_guards_in_lifo_order() {
+        let initial = Hub::current();
+        let hub_a = Arc::new(Hub::new_from_top(initial.clone()));
+        let hub_b = Arc::new(Hub::new_from_top(hub_a.clone()));
+
+        assert!(!Arc::ptr_eq(&hub_b, &initial));
+        assert!(!Arc::ptr_eq(&hub_a, &initial));
+        assert!(!Arc::ptr_eq(&hub_a, &hub_b));
+
+        let mut stack = SpanGuardStack::new();
+        let id = SpanId::from_u64(2);
+
+        stack.push(id.clone(), HubSwitchGuard::new(hub_a.clone()));
+        assert!(Arc::ptr_eq(&Hub::current(), &hub_a));
+
+        stack.push(id.clone(), HubSwitchGuard::new(hub_b.clone()));
+        assert!(Arc::ptr_eq(&Hub::current(), &hub_b));
+
+        stack.remove(&id);
+        assert!(Arc::ptr_eq(&Hub::current(), &initial));
+    }
+}

sentry-tracing/tests/span_cross_thread.rs

@@ -0,0 +1,77 @@
+mod shared;
+
+use sentry::protocol::Context;
+use std::thread;
+use std::time::Duration;
+
+#[test]
+fn cross_thread_span_entries_share_transaction() {
+    let transport = shared::init_sentry(1.0);
+
+    let span = tracing::info_span!("foo");
+    let span2 = span.clone();
+
+    let handle1 = thread::spawn(move || {
+        let _guard = span.enter();
+        let _bar_span = tracing::info_span!("bar").entered();
+        thread::sleep(Duration::from_millis(100));
+    });
+
+    let handle2 = thread::spawn(move || {
+        thread::sleep(Duration::from_millis(10));
+        let _guard = span2.enter();
+        let _baz_span = tracing::info_span!("baz").entered();
+        thread::sleep(Duration::from_millis(50));
+    });
+
+    handle1.join().unwrap();
+    handle2.join().unwrap();
+
+    let data = transport.fetch_and_clear_envelopes();
+    let transactions: Vec<_> = data
+        .into_iter()
+        .flat_map(|envelope| {
+            envelope
+                .items()
+                .filter_map(|item| match item {
+                    sentry::protocol::EnvelopeItem::Transaction(transaction) => {
+                        Some(transaction.clone())
+                    }
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+        })
+        .collect();
+
+    assert_eq!(
+        transactions.len(),
+        1,
+        "expected a single transaction for cross-thread span entries"
+    );
+
+    let transaction = &transactions[0];
+    assert_eq!(transaction.name.as_deref(), Some("foo"));
+
+    let trace = match transaction
+        .contexts
+        .get("trace")
+        .expect("transaction should include trace context")
+    {
+        Context::Trace(trace) => trace,
+        unexpected => panic!("expected trace context but got {unexpected:?}"),
+    };
+
+    let bar_span = transaction
+        .spans
+        .iter()
+        .find(|span| span.description.as_deref() == Some("bar"))
+        .expect("expected span \"bar\" to be recorded in the transaction");
+    let baz_span = transaction
+        .spans
+        .iter()
+        .find(|span| span.description.as_deref() == Some("baz"))
+        .expect("expected span \"baz\" to be recorded in the transaction");
+
+    assert_eq!(bar_span.parent_span_id, Some(trace.span_id));
+    assert_eq!(baz_span.parent_span_id, Some(trace.span_id));
+}

sentry-tracing/tests/span_reentrancy.rs

@@ -0,0 +1,70 @@
+use sentry::protocol::EnvelopeItem;
+
+mod shared;
+
+/// Ensures re-entering the same span does not corrupt the current tracing state,
+/// so subsequent spans are still recorded under a single transaction.
+#[test]
+fn reentering_span_preserves_parent() {
+    let transport = shared::init_sentry(1.0);
+
+    {
+        // Create a span and enter it, then re-enter the same span to simulate
+        // common async polling behavior where a span can be entered multiple times.
+        let span_a = tracing::info_span!("a");
+        let _enter_a = span_a.enter();
+        {
+            let _reenter_a = span_a.enter();
+        }
+
+        // Create another span while the original span is still entered to ensure
+        // it is recorded on the same transaction rather than starting a new one.
+        let span_b = tracing::info_span!("b");
+        {
+            let _enter_b = span_b.enter();
+        }
+    }
+
+    let transactions: Vec<_> = transport
+        .fetch_and_clear_envelopes()
+        .into_iter()
+        .flat_map(|envelope| envelope.into_items())
+        .filter_map(|item| match item {
+            EnvelopeItem::Transaction(transaction) => Some(transaction),
+            _ => None,
+        })
+        .collect();
+
+    assert_eq!(
+        transactions.len(),
+        1,
+        "expected a single transaction when reentering a span"
+    );
+
+    let transaction = &transactions[0];
+    assert_eq!(transaction.name.as_deref(), Some("a"));
+
+    let trace = match transaction
+        .contexts
+        .get("trace")
+        .expect("transaction should include trace context")
+    {
+        sentry::protocol::Context::Trace(trace) => trace,
+        unexpected => panic!("expected trace context but got {unexpected:?}"),
+    };
+
+    let b_span = transaction
+        .spans
+        .iter()
+        .find(|span| span.description.as_deref() == Some("b"))
+        .expect("expected span \"b\" to be recorded in the transaction");
+
+    assert_eq!(b_span.parent_span_id, Some(trace.span_id));
+    assert!(
+        !transaction
+            .spans
+            .iter()
+            .any(|span| span.description.as_deref() == Some("a")),
+        "expected the transaction root span not to be duplicated in span list"
+    );
+}