Skip to content

Commit

Permalink
fix(sentry-tracing): switch sentry spans on enter and exit
Browse files Browse the repository at this point in the history
  • Loading branch information
saiintbrisson committed Sep 27, 2024
1 parent 96b35c7 commit a138e3c
Show file tree
Hide file tree
Showing 4 changed files with 101 additions and 21 deletions.
26 changes: 13 additions & 13 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

42 changes: 42 additions & 0 deletions sentry-core/src/performance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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") {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
48 changes: 40 additions & 8 deletions sentry-tracing/src/layer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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>) {
Expand All @@ -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
Expand Down
6 changes: 6 additions & 0 deletions sentry-types/src/protocol/v7.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down

0 comments on commit a138e3c

Please sign in to comment.