Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix rust tracing #3817

Merged
merged 1 commit into from
Nov 28, 2024
Merged

Fix rust tracing #3817

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 27 additions & 42 deletions sentry_sdk/integrations/rust_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,9 @@
import sentry_sdk
from sentry_sdk.integrations import Integration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import Span as SentrySpan
from sentry_sdk.tracing import POTelSpan as SentrySpan
from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE

TraceState = Optional[Tuple[Optional[SentrySpan], SentrySpan]]


class RustTracingLevel(Enum):
Trace: str = "TRACE"
Expand Down Expand Up @@ -171,7 +169,7 @@ def _include_tracing_fields(self) -> bool:
else self.include_tracing_fields
)

def on_event(self, event: str, _span_state: TraceState) -> None:
def on_event(self, event: str, _span_state: Optional[SentrySpan]) -> None:
deserialized_event = json.loads(event)
metadata = deserialized_event.get("metadata", {})

Expand All @@ -185,7 +183,7 @@ def on_event(self, event: str, _span_state: TraceState) -> None:
elif event_type == EventTypeMapping.Event:
process_event(deserialized_event)

def on_new_span(self, attrs: str, span_id: str) -> TraceState:
def on_new_span(self, attrs: str, span_id: str) -> Optional[SentrySpan]:
attrs = json.loads(attrs)
metadata = attrs.get("metadata", {})

Expand All @@ -205,48 +203,35 @@ def on_new_span(self, attrs: str, span_id: str) -> TraceState:
else:
sentry_span_name = "<unknown>"

kwargs = {
"op": "function",
"name": sentry_span_name,
"origin": self.origin,
}

scope = sentry_sdk.get_current_scope()
parent_sentry_span = scope.span
if parent_sentry_span:
sentry_span = parent_sentry_span.start_child(**kwargs)
else:
sentry_span = scope.start_span(**kwargs)
span = sentry_sdk.start_span(
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@matt-codecov @antonpirker I simplified this part since we don't want users to do scope stuff in integrations, it is anti-patterny

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

from a quick look at the code, it looks like __enter__() will still update the current span and keep track of the parent so that __exit__() can restore it, right? if so then i am fine with this. not qualified to stamp, and also CI is unhappy

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes exactly, span management should be done in the core and will just work
ci is unhappy because this is a dev branch

op="function",
name=sentry_span_name,
origin=self.origin,
only_if_parent=True,
)
span.__enter__()

fields = metadata.get("fields", [])
for field in fields:
if self._include_tracing_fields():
sentry_span.set_data(field, attrs.get(field))
else:
sentry_span.set_data(field, SENSITIVE_DATA_SUBSTITUTE)

scope.span = sentry_span
return (parent_sentry_span, sentry_span)

def on_close(self, span_id: str, span_state: TraceState) -> None:
if span_state is None:
return

parent_sentry_span, sentry_span = span_state
sentry_span.finish()
sentry_sdk.get_current_scope().span = parent_sentry_span

def on_record(self, span_id: str, values: str, span_state: TraceState) -> None:
if span_state is None:
return
_parent_sentry_span, sentry_span = span_state

deserialized_values = json.loads(values)
for key, value in deserialized_values.items():
if self._include_tracing_fields():
sentry_span.set_data(key, value)
span.set_data(field, attrs.get(field))
else:
sentry_span.set_data(key, SENSITIVE_DATA_SUBSTITUTE)
span.set_data(field, SENSITIVE_DATA_SUBSTITUTE)

return span

def on_close(self, span_id: str, span: Optional[SentrySpan]) -> None:
if span is not None:
span.__exit__(None, None, None)

def on_record(self, span_id: str, values: str, span: Optional[SentrySpan]) -> None:
if span is not None:
deserialized_values = json.loads(values)
for key, value in deserialized_values.items():
if self._include_tracing_fields():
span.set_data(key, value)
else:
span.set_data(key, SENSITIVE_DATA_SUBSTITUTE)


class RustTracingIntegration(Integration):
Expand Down
9 changes: 7 additions & 2 deletions sentry_sdk/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
import uuid
import random
import time
Expand Down Expand Up @@ -1631,8 +1632,12 @@ def finish(self, end_timestamp=None):

def to_json(self):
# type: () -> dict[str, Any]
# TODO-neel-potel for sampling context
pass
"""
Only meant for testing. Not used internally anymore.
"""
if not isinstance(self._otel_span, ReadableSpan):
return {}
return json.loads(self._otel_span.to_json())

def get_trace_context(self):
# type: () -> dict[str, Any]
Expand Down
82 changes: 42 additions & 40 deletions tests/integrations/rust_tracing/test_rust_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@
RustTracingLevel,
EventTypeMapping,
)
from sentry_sdk import start_transaction, capture_message
from sentry_sdk import start_span, capture_message
from tests.conftest import ApproxDict


def _test_event_type_mapping(metadata: Dict[str, object]) -> EventTypeMapping:
Expand Down Expand Up @@ -74,11 +75,11 @@ def test_on_new_span_on_close(sentry_init, capture_events):
sentry_init(integrations=[integration], traces_sample_rate=1.0)

events = capture_events()
with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

sentry_first_rust_span = sentry_sdk.get_current_span()
_, rust_first_rust_span = rust_tracing.spans[3]
rust_first_rust_span = rust_tracing.spans[3]

assert sentry_first_rust_span == rust_first_rust_span

Expand All @@ -102,7 +103,7 @@ def test_on_new_span_on_close(sentry_init, capture_events):
data = span["data"]
assert data["use_memoized"]
assert data["index"] == 10
assert data["version"] is None
assert "version" not in data


def test_nested_on_new_span_on_close(sentry_init, capture_events):
Expand All @@ -115,23 +116,20 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events):
sentry_init(integrations=[integration], traces_sample_rate=1.0)

events = capture_events()
with start_transaction():
with start_span():
original_sentry_span = sentry_sdk.get_current_span()

rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10)
sentry_first_rust_span = sentry_sdk.get_current_span()
_, rust_first_rust_span = rust_tracing.spans[3]
rust_first_rust_span = rust_tracing.spans[3]

# Use a different `index_arg` value for the inner span to help
# distinguish the two at the end of the test
rust_tracing.new_span(RustTracingLevel.Info, 5, index_arg=9)
sentry_second_rust_span = sentry_sdk.get_current_span()
rust_parent_span, rust_second_rust_span = rust_tracing.spans[5]
rust_second_rust_span = rust_tracing.spans[5]

assert rust_second_rust_span == sentry_second_rust_span
assert rust_parent_span == sentry_first_rust_span
assert rust_parent_span == rust_first_rust_span
assert rust_parent_span != rust_second_rust_span

rust_tracing.close_span(5)

Expand Down Expand Up @@ -171,12 +169,12 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events):
first_span_data = first_span["data"]
assert first_span_data["use_memoized"]
assert first_span_data["index"] == 10
assert first_span_data["version"] is None
assert "version" not in first_span_data

second_span_data = second_span["data"]
assert second_span_data["use_memoized"]
assert second_span_data["index"] == 9
assert second_span_data["version"] is None
assert "version" not in second_span_data


def test_on_new_span_without_transaction(sentry_init):
Expand Down Expand Up @@ -207,7 +205,7 @@ def test_on_event_exception(sentry_init, capture_events):
events = capture_events()
sentry_sdk.get_isolation_scope().clear_breadcrumbs()

with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

# Mapped to Exception
Expand Down Expand Up @@ -243,7 +241,7 @@ def test_on_event_breadcrumb(sentry_init, capture_events):
events = capture_events()
sentry_sdk.get_isolation_scope().clear_breadcrumbs()

with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

# Mapped to Breadcrumb
Expand Down Expand Up @@ -274,7 +272,7 @@ def test_on_event_event(sentry_init, capture_events):
events = capture_events()
sentry_sdk.get_isolation_scope().clear_breadcrumbs()

with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

# Mapped to Event
Expand Down Expand Up @@ -311,7 +309,7 @@ def test_on_event_ignored(sentry_init, capture_events):
events = capture_events()
sentry_sdk.get_isolation_scope().clear_breadcrumbs()

with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

# Ignored
Expand Down Expand Up @@ -344,7 +342,7 @@ def span_filter(metadata: Dict[str, object]) -> bool:
sentry_init(integrations=[integration], traces_sample_rate=1.0)

events = capture_events()
with start_transaction():
with start_span():
original_sentry_span = sentry_sdk.get_current_span()

# Span is not ignored
Expand Down Expand Up @@ -377,16 +375,16 @@ def test_record(sentry_init):
)
sentry_init(integrations=[integration], traces_sample_rate=1.0)

with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

span_before_record = sentry_sdk.get_current_span().to_json()
assert span_before_record["data"]["version"] is None
assert "version" not in span_before_record["attributes"]

rust_tracing.record(3)

span_after_record = sentry_sdk.get_current_span().to_json()
assert span_after_record["data"]["version"] == "memoized"
assert span_after_record["attributes"]["version"] == "memoized"


def test_record_in_ignored_span(sentry_init):
Expand All @@ -403,18 +401,18 @@ def span_filter(metadata: Dict[str, object]) -> bool:
)
sentry_init(integrations=[integration], traces_sample_rate=1.0)

with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

span_before_record = sentry_sdk.get_current_span().to_json()
assert span_before_record["data"]["version"] is None
assert "version" not in span_before_record["attributes"]

rust_tracing.new_span(RustTracingLevel.Trace, 5)
rust_tracing.record(5)

# `on_record()` should not do anything to the current Sentry span if the associated Rust span was ignored
span_after_record = sentry_sdk.get_current_span().to_json()
assert span_after_record["data"]["version"] is None
assert "version" not in span_after_record["attributes"]


@pytest.mark.parametrize(
Expand Down Expand Up @@ -443,33 +441,37 @@ def test_include_tracing_fields(
traces_sample_rate=1.0,
send_default_pii=send_default_pii,
)
with start_transaction():
with start_span():
rust_tracing.new_span(RustTracingLevel.Info, 3)

span_before_record = sentry_sdk.get_current_span().to_json()
if tracing_fields_expected:
assert span_before_record["data"]["version"] is None
assert "version" not in span_before_record["attributes"]
else:
assert span_before_record["data"]["version"] == "[Filtered]"
assert span_before_record["attributes"]["version"] == "[Filtered]"

rust_tracing.record(3)

span_after_record = sentry_sdk.get_current_span().to_json()

if tracing_fields_expected:
assert span_after_record["data"] == {
"thread.id": mock.ANY,
"thread.name": mock.ANY,
"use_memoized": True,
"version": "memoized",
"index": 10,
}
assert span_after_record["attributes"] == ApproxDict(
{
"thread.id": mock.ANY,
"thread.name": mock.ANY,
"use_memoized": True,
"version": "memoized",
"index": 10,
}
)

else:
assert span_after_record["data"] == {
"thread.id": mock.ANY,
"thread.name": mock.ANY,
"use_memoized": "[Filtered]",
"version": "[Filtered]",
"index": "[Filtered]",
}
assert span_after_record["attributes"] == ApproxDict(
{
"thread.id": mock.ANY,
"thread.name": mock.ANY,
"use_memoized": "[Filtered]",
"version": "[Filtered]",
"index": "[Filtered]",
}
)
Loading