Skip to content

Commit 0ff374f

Browse files
lvthanh03george-sentry
authored andcommitted
feat(spans): Skip enrichment when message contains skip_enrichment=True flag (#112024)
Going off of product requirements from #111820 (comment), we want to have a way to signal the process-segments consumer to skip enrichment once a segment hits the size limit. When this happens, the process-spans consumer would produce the segment in chunks to the buffered-segments topic, where each message contains one chunk along with the flag `skip_enrichment=True`.
1 parent da938c9 commit 0ff374f

File tree

3 files changed

+49
-4
lines changed

3 files changed

+49
-4
lines changed

src/sentry/spans/consumers/process_segments/factory.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,10 @@ def _process_message(
115115
try:
116116
value = message.payload.value
117117
segment = orjson.loads(value)
118-
processed = process_segment(segment["spans"], skip_produce=skip_produce)
118+
skip_enrichment = segment.get("skip_enrichment", False)
119+
processed = process_segment(
120+
segment["spans"], skip_produce=skip_produce, skip_enrichment=skip_enrichment
121+
)
119122
return [_serialize_payload(span, message.timestamp) for span in processed]
120123
except Exception:
121124
logger.exception("segments.invalid-message")

src/sentry/spans/consumers/process_segments/message.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050

5151
@metrics.wraps("spans.consumers.process_segments.process_segment")
5252
def process_segment(
53-
unprocessed_spans: list[SpanEvent], skip_produce: bool = False
53+
unprocessed_spans: list[SpanEvent], skip_produce: bool = False, skip_enrichment: bool = False
5454
) -> list[CompatibleSpan]:
5555
sample_rate = (
5656
settings.SENTRY_PROCESS_SEGMENTS_TRANSACTIONS_SAMPLE_RATE
@@ -62,14 +62,17 @@ def process_segment(
6262
"sample_rate": sample_rate,
6363
},
6464
):
65-
return _process_segment(unprocessed_spans, skip_produce)
65+
return _process_segment(unprocessed_spans, skip_produce, skip_enrichment)
6666

6767

6868
def _process_segment(
69-
unprocessed_spans: list[SpanEvent], skip_produce: bool
69+
unprocessed_spans: list[SpanEvent], skip_produce: bool, skip_enrichment: bool
7070
) -> list[CompatibleSpan]:
7171
_verify_compatibility(unprocessed_spans)
7272

73+
if skip_enrichment:
74+
return [make_compatible(span) for span in unprocessed_spans]
75+
7376
if unprocessed_spans:
7477
project_id = unprocessed_spans[0].get("project_id")
7578
if project_id in options.get("spans.process-segments.skip-enrichment-projects"):

tests/sentry/spans/consumers/process_segments/test_message.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -405,6 +405,45 @@ def test_no_skip_enrichment_for_other_project(self, mock_enrich: mock.MagicMock)
405405

406406
mock_enrich.assert_called_once()
407407

408+
@mock.patch(
409+
"sentry.spans.consumers.process_segments.message.TreeEnricher.enrich_spans",
410+
wraps=None,
411+
)
412+
def test_skip_enrichment_flag(self, mock_enrich: mock.MagicMock) -> None:
413+
"""Test that enrichment is skipped when skip_enrichment=True is passed."""
414+
segment_span = build_mock_span(
415+
project_id=self.project.id,
416+
is_segment=True,
417+
)
418+
child_span = build_mock_span(
419+
project_id=self.project.id,
420+
parent_span_id=segment_span["span_id"],
421+
)
422+
423+
processed_spans = process_segment([child_span, segment_span], skip_enrichment=True)
424+
425+
mock_enrich.assert_not_called()
426+
assert len(processed_spans) == 2
427+
428+
@mock.patch(
429+
"sentry.spans.consumers.process_segments.message.TreeEnricher.enrich_spans",
430+
)
431+
def test_no_skip_enrichment_when_flag_is_false(self, mock_enrich: mock.MagicMock) -> None:
432+
"""Test that enrichment runs normally when skip_enrichment=False."""
433+
mock_enrich.return_value = (None, [])
434+
segment_span = build_mock_span(
435+
project_id=self.project.id,
436+
is_segment=True,
437+
)
438+
child_span = build_mock_span(
439+
project_id=self.project.id,
440+
parent_span_id=segment_span["span_id"],
441+
)
442+
443+
process_segment([child_span, segment_span], skip_enrichment=False)
444+
445+
mock_enrich.assert_called_once()
446+
408447

409448
@exclude_experimental_detectors
410449
class TestSegmentDropKillswitch(TestCase):

0 commit comments

Comments
 (0)