From 4d899e79cd1e1cd33e8e4f0cd5d0dc8824e6b4fb Mon Sep 17 00:00:00 2001 From: Tony Le Date: Fri, 3 Jul 2026 07:19:58 -0400 Subject: [PATCH 1/7] feat(spans): add option to rollout process_segment task --- src/sentry/options/defaults.py | 6 + src/sentry/spans/consumers/process/flusher.py | 18 ++- .../spans/consumers/process/test_flusher.py | 131 ++++++++++++++++-- 3 files changed, 135 insertions(+), 20 deletions(-) diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py index 7d8017777a31..e99d292716ba 100644 --- a/src/sentry/options/defaults.py +++ b/src/sentry/options/defaults.py @@ -2977,6 +2977,12 @@ default=False, flags=FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE, ) +register( + "spans.buffer.process-segments-task-rollout-rate", + type=Float, + default=0.0, + flags=FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE, +) # List of trace_ids to enable debug logging for. Empty = debug off. # When set, logs detailed metrics about zunionstore set sizes, key existence, and trace structure. diff --git a/src/sentry/spans/consumers/process/flusher.py b/src/sentry/spans/consumers/process/flusher.py index 47d63d012147..abfcc579e64e 100644 --- a/src/sentry/spans/consumers/process/flusher.py +++ b/src/sentry/spans/consumers/process/flusher.py @@ -21,14 +21,17 @@ from sentry.conf.types.kafka_definition import Topic from sentry.constants import DataCategory from sentry.models.project import Project +from sentry.options.rollout import in_random_rollout from sentry.processing.backpressure.memory import ServiceMemory from sentry.spans.buffer import SpansBuffer +from sentry.spans.consumers.process_segments.tasks import process_segment_task from sentry.utils import metrics from sentry.utils.arroyo import run_with_initialized_sentry from sentry.utils.kafka_config import get_kafka_producer_cluster_options, get_topic_definition from sentry.utils.outcomes import Outcome, track_outcome MAX_PROCESS_RESTARTS = 10 +PROCESS_SEGMENTS_TASK_ROLLOUT_OPTION = "spans.buffer.process-segments-task-rollout-rate" logger = logging.getLogger(__name__) @@ -358,11 +361,16 @@ def produce(project_id: int, payload: KafkaPayload, dropped: int) -> None: len(kafka_payload.value), tags={"shard": shard_tag}, ) - produce( - flushed_segment.project_id, - kafka_payload, - len(message["spans"]), - ) + if produce_to_pipe is None and in_random_rollout( + PROCESS_SEGMENTS_TASK_ROLLOUT_OPTION + ): + process_segment_task.apply_async(args=[kafka_payload.value]) + else: + produce( + flushed_segment.project_id, + kafka_payload, + len(message["spans"]), + ) with metrics.timer("spans.buffer.flusher.wait_produce", tags={"shards": shard_tag}): for project_id, future, dropped in producer_futures: diff --git a/tests/sentry/spans/consumers/process/test_flusher.py b/tests/sentry/spans/consumers/process/test_flusher.py index 545bbe9b3db2..3be4203f32e9 100644 --- a/tests/sentry/spans/consumers/process/test_flusher.py +++ b/tests/sentry/spans/consumers/process/test_flusher.py @@ -1,4 +1,5 @@ import time +from concurrent.futures import Future from time import sleep from types import SimpleNamespace from typing import Any @@ -6,6 +7,7 @@ import orjson import pytest +from arroyo.backends.kafka import KafkaPayload from arroyo.processing.strategies.noop import Noop from django.test import override_settings @@ -21,6 +23,32 @@ def _payload(span_id: str) -> bytes: return orjson.dumps({"span_id": span_id}) +def _buffer_with_segment( + *, + project_id: int = 999_002, + trace_id: str = "9" * 32, + span_id: str = "b" * 16, + slice_id: int = 999_002, +) -> SpansBuffer: + buffer = SpansBuffer(assigned_shards=[0], slice_id=slice_id) + buffer.process_spans( + [ + Span( + payload=_payload(span_id), + trace_id=trace_id, + span_id=span_id, + parent_span_id=None, + segment_id=None, + is_segment_span=True, + project_id=project_id, + partition=0, + ) + ], + now=0, + ) + return buffer + + def _blocking_main_for_join_test( buffer, shards, stopped, current_drift, backpressure_since, healthy_since, produce_to_pipe ): @@ -39,21 +67,8 @@ def test_flusher_logs_flushed_segments() -> None: slice_id = 999_002 segment_key = f"span-buf:s:{{{project_id}:{trace_id}}}:{span_id}".encode() queue_key = f"span-buf:q:{slice_id}-0".encode() - buffer = SpansBuffer(assigned_shards=[0], slice_id=slice_id) - buffer.process_spans( - [ - Span( - payload=_payload(span_id), - trace_id=trace_id, - span_id=span_id, - parent_span_id=None, - segment_id=None, - is_segment_span=True, - project_id=project_id, - partition=0, - ) - ], - now=0, + buffer = _buffer_with_segment( + project_id=project_id, trace_id=trace_id, span_id=span_id, slice_id=slice_id ) stopped = SimpleNamespace(value=0) current_drift = SimpleNamespace(value=0) @@ -93,6 +108,92 @@ def produce_to_pipe(project_id: int, payload: Any, dropped: int) -> None: assert not buffer.client.keys(f"*{project_id}:{trace_id}*") +@override_options({**DEFAULT_OPTIONS, "spans.buffer.process-segments-task-rollout-rate": 0.0}) +def test_flusher_produces_flushed_segments_to_buffered_segments_topic() -> None: + span_id = "b" * 16 + buffer = _buffer_with_segment(span_id=span_id) + stopped = SimpleNamespace(value=0) + current_drift = SimpleNamespace(value=0) + backpressure_since = SimpleNamespace(value=0) + healthy_since = SimpleNamespace(value=0) + produced_payloads: list[KafkaPayload] = [] + producer_future: Future[None] = Future() + producer_future.set_result(None) + producer_manager = mock.Mock() + + def produce(payload: KafkaPayload) -> Future[None]: + produced_payloads.append(payload) + stopped.value = 1 + return producer_future + + producer_manager.produce.side_effect = produce + + with ( + mock.patch( + "sentry.spans.consumers.process.flusher.MultiProducer", + return_value=producer_manager, + ) as mock_multi_producer, + mock.patch( + "sentry.spans.consumers.process.flusher.process_segment_task.apply_async" + ) as mock_apply_async, + ): + SpanFlusher.main( + buffer, + shards=[0], + stopped=stopped, + current_drift=current_drift, + backpressure_since=backpressure_since, + healthy_since=healthy_since, + produce_to_pipe=None, + ) + + mock_multi_producer.assert_called_once_with(Topic.BUFFERED_SEGMENTS) + mock_apply_async.assert_not_called() + assert len(produced_payloads) == 1 + assert orjson.loads(produced_payloads[0].value)["spans"][0]["span_id"] == span_id + + +@override_options({**DEFAULT_OPTIONS, "spans.buffer.process-segments-task-rollout-rate": 1.0}) +def test_flusher_produces_flushed_segments_to_process_segment_task() -> None: + span_id = "c" * 16 + buffer = _buffer_with_segment(span_id=span_id) + stopped = SimpleNamespace(value=0) + current_drift = SimpleNamespace(value=0) + backpressure_since = SimpleNamespace(value=0) + healthy_since = SimpleNamespace(value=0) + producer_manager = mock.Mock() + + def apply_async(*args: Any, **kwargs: Any) -> None: + stopped.value = 1 + + with ( + mock.patch( + "sentry.spans.consumers.process.flusher.MultiProducer", + return_value=producer_manager, + ) as mock_multi_producer, + mock.patch( + "sentry.spans.consumers.process.flusher.process_segment_task.apply_async", + side_effect=apply_async, + ) as mock_apply_async, + ): + SpanFlusher.main( + buffer, + shards=[0], + stopped=stopped, + current_drift=current_drift, + backpressure_since=backpressure_since, + healthy_since=healthy_since, + produce_to_pipe=None, + ) + + mock_multi_producer.assert_called_once_with(Topic.BUFFERED_SEGMENTS) + producer_manager.produce.assert_not_called() + mock_apply_async.assert_called_once() + task_args = mock_apply_async.call_args.kwargs["args"] + assert len(task_args) == 1 + assert orjson.loads(task_args[0])["spans"][0]["span_id"] == span_id + + @override_options({**DEFAULT_OPTIONS, "spans.buffer.max-flush-segments": 1}) def test_backpressure() -> None: # Flush very aggressively to make join() faster From 6579d20dab282e797f3ebbd907bdf665d66e15e2 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Fri, 3 Jul 2026 10:08:10 -0400 Subject: [PATCH 2/7] add wait_for_delievry for the process_segment task --- src/sentry/spans/consumers/process_segments/tasks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/sentry/spans/consumers/process_segments/tasks.py b/src/sentry/spans/consumers/process_segments/tasks.py index 5fa6473fcb67..d450125de162 100644 --- a/src/sentry/spans/consumers/process_segments/tasks.py +++ b/src/sentry/spans/consumers/process_segments/tasks.py @@ -37,6 +37,7 @@ def _get_snuba_items_producer( name="sentry.spans.process_segments.process_segment", namespace=spans_process_segments_tasks, at_most_once=True, + wait_for_delivery=True, compression_type=CompressionType.ZSTD, silo_mode=SiloMode.CELL, ) From ac6bbf65aa717742cbefbe21b7e8f6c86879aed2 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Mon, 6 Jul 2026 13:30:36 -0400 Subject: [PATCH 3/7] use apply_async_with_future --- src/sentry/spans/consumers/process/flusher.py | 12 ++++++- .../spans/consumers/process_segments/tasks.py | 1 - src/sentry/taskworker/silolimiter.py | 2 +- .../spans/consumers/process/test_flusher.py | 36 +++++++++++++------ tests/sentry/tasks/test_base.py | 2 +- 5 files changed, 39 insertions(+), 14 deletions(-) diff --git a/src/sentry/spans/consumers/process/flusher.py b/src/sentry/spans/consumers/process/flusher.py index abfcc579e64e..c0c00a7660a7 100644 --- a/src/sentry/spans/consumers/process/flusher.py +++ b/src/sentry/spans/consumers/process/flusher.py @@ -364,7 +364,17 @@ def produce(project_id: int, payload: KafkaPayload, dropped: int) -> None: if produce_to_pipe is None and in_random_rollout( PROCESS_SEGMENTS_TASK_ROLLOUT_OPTION ): - process_segment_task.apply_async(args=[kafka_payload.value]) + task_produce_future = process_segment_task.apply_async_with_future( + args=[kafka_payload.value] + ) + if task_produce_future is not None: + producer_futures.append( + ( + flushed_segment.project_id, + task_produce_future, + len(message["spans"]), + ) + ) else: produce( flushed_segment.project_id, diff --git a/src/sentry/spans/consumers/process_segments/tasks.py b/src/sentry/spans/consumers/process_segments/tasks.py index d450125de162..5fa6473fcb67 100644 --- a/src/sentry/spans/consumers/process_segments/tasks.py +++ b/src/sentry/spans/consumers/process_segments/tasks.py @@ -37,7 +37,6 @@ def _get_snuba_items_producer( name="sentry.spans.process_segments.process_segment", namespace=spans_process_segments_tasks, at_most_once=True, - wait_for_delivery=True, compression_type=CompressionType.ZSTD, silo_mode=SiloMode.CELL, ) diff --git a/src/sentry/taskworker/silolimiter.py b/src/sentry/taskworker/silolimiter.py index 6b3ec27689c8..e24bdd915224 100644 --- a/src/sentry/taskworker/silolimiter.py +++ b/src/sentry/taskworker/silolimiter.py @@ -30,7 +30,7 @@ def handle(*args: P.args, **kwargs: P.kwargs) -> Any: def __call__(self, decorated_task: Task[P, R]) -> Task[P, R]: # Replace the sentry.taskworker.Task interface used to schedule tasks. - replacements = {"delay", "apply_async"} + replacements = {"delay", "apply_async", "apply_async_with_future"} for attr_name in replacements: task_attr = getattr(decorated_task, attr_name) if callable(task_attr): diff --git a/tests/sentry/spans/consumers/process/test_flusher.py b/tests/sentry/spans/consumers/process/test_flusher.py index 3be4203f32e9..257b5e9ed2c7 100644 --- a/tests/sentry/spans/consumers/process/test_flusher.py +++ b/tests/sentry/spans/consumers/process/test_flusher.py @@ -134,8 +134,8 @@ def produce(payload: KafkaPayload) -> Future[None]: return_value=producer_manager, ) as mock_multi_producer, mock.patch( - "sentry.spans.consumers.process.flusher.process_segment_task.apply_async" - ) as mock_apply_async, + "sentry.spans.consumers.process.flusher.process_segment_task.apply_async_with_future" + ) as mock_apply_async_with_future, ): SpanFlusher.main( buffer, @@ -148,23 +148,37 @@ def produce(payload: KafkaPayload) -> Future[None]: ) mock_multi_producer.assert_called_once_with(Topic.BUFFERED_SEGMENTS) - mock_apply_async.assert_not_called() + mock_apply_async_with_future.assert_not_called() assert len(produced_payloads) == 1 assert orjson.loads(produced_payloads[0].value)["spans"][0]["span_id"] == span_id @override_options({**DEFAULT_OPTIONS, "spans.buffer.process-segments-task-rollout-rate": 1.0}) def test_flusher_produces_flushed_segments_to_process_segment_task() -> None: + project_id = 999_002 + trace_id = "9" * 32 span_id = "c" * 16 - buffer = _buffer_with_segment(span_id=span_id) + slice_id = 999_002 + segment_key = f"span-buf:s:{{{project_id}:{trace_id}}}:{span_id}".encode() + queue_key = f"span-buf:q:{slice_id}-0".encode() + buffer = _buffer_with_segment( + project_id=project_id, trace_id=trace_id, span_id=span_id, slice_id=slice_id + ) stopped = SimpleNamespace(value=0) current_drift = SimpleNamespace(value=0) backpressure_since = SimpleNamespace(value=0) healthy_since = SimpleNamespace(value=0) producer_manager = mock.Mock() + task_future = mock.Mock() - def apply_async(*args: Any, **kwargs: Any) -> None: + def wait_for_task_delivery() -> None: + assert buffer.client.zscore(queue_key, segment_key) is not None + + task_future.result.side_effect = wait_for_task_delivery + + def apply_async_with_future(*args: Any, **kwargs: Any) -> Any: stopped.value = 1 + return task_future with ( mock.patch( @@ -172,9 +186,9 @@ def apply_async(*args: Any, **kwargs: Any) -> None: return_value=producer_manager, ) as mock_multi_producer, mock.patch( - "sentry.spans.consumers.process.flusher.process_segment_task.apply_async", - side_effect=apply_async, - ) as mock_apply_async, + "sentry.spans.consumers.process.flusher.process_segment_task.apply_async_with_future", + side_effect=apply_async_with_future, + ) as mock_apply_async_with_future, ): SpanFlusher.main( buffer, @@ -188,10 +202,12 @@ def apply_async(*args: Any, **kwargs: Any) -> None: mock_multi_producer.assert_called_once_with(Topic.BUFFERED_SEGMENTS) producer_manager.produce.assert_not_called() - mock_apply_async.assert_called_once() - task_args = mock_apply_async.call_args.kwargs["args"] + mock_apply_async_with_future.assert_called_once() + task_args = mock_apply_async_with_future.call_args.kwargs["args"] assert len(task_args) == 1 assert orjson.loads(task_args[0])["spans"][0]["span_id"] == span_id + task_future.result.assert_called_once_with() + assert buffer.client.zscore(queue_key, segment_key) is None @override_options({**DEFAULT_OPTIONS, "spans.buffer.max-flush-segments": 1}) diff --git a/tests/sentry/tasks/test_base.py b/tests/sentry/tasks/test_base.py index 1cf22fbf3202..30441fbab3e9 100644 --- a/tests/sentry/tasks/test_base.py +++ b/tests/sentry/tasks/test_base.py @@ -93,7 +93,7 @@ def test_task_silo_limit_call_monolith() -> None: @pytest.mark.parametrize( "method_name", - ("apply_async", "delay"), + ("apply_async", "apply_async_with_future", "delay"), ) @override_settings(SILO_MODE=SiloMode.CONTROL) def test_task_silo_limit_task_methods(method_name: str) -> None: From f4f5e663b5342dd4b6a3a987e54795e6d6dbcf03 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Mon, 6 Jul 2026 13:31:11 -0400 Subject: [PATCH 4/7] use spans app feature --- src/sentry/taskworker/namespaces.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sentry/taskworker/namespaces.py b/src/sentry/taskworker/namespaces.py index 9d5fca72f59c..b368f913a944 100644 --- a/src/sentry/taskworker/namespaces.py +++ b/src/sentry/taskworker/namespaces.py @@ -94,7 +94,7 @@ spans_process_segments_tasks = app.taskregistry.create_namespace( "spans.process_segments", - app_feature="transactions", + app_feature="spans", ) ingest_attachments_tasks = app.taskregistry.create_namespace( From b8d9053146f40215e09d1c1426dc47b218f6fc21 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Mon, 6 Jul 2026 13:40:30 -0400 Subject: [PATCH 5/7] use the option directly --- src/sentry/spans/consumers/process/flusher.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/sentry/spans/consumers/process/flusher.py b/src/sentry/spans/consumers/process/flusher.py index c0c00a7660a7..bc2bdd673276 100644 --- a/src/sentry/spans/consumers/process/flusher.py +++ b/src/sentry/spans/consumers/process/flusher.py @@ -31,7 +31,6 @@ from sentry.utils.outcomes import Outcome, track_outcome MAX_PROCESS_RESTARTS = 10 -PROCESS_SEGMENTS_TASK_ROLLOUT_OPTION = "spans.buffer.process-segments-task-rollout-rate" logger = logging.getLogger(__name__) @@ -362,7 +361,7 @@ def produce(project_id: int, payload: KafkaPayload, dropped: int) -> None: tags={"shard": shard_tag}, ) if produce_to_pipe is None and in_random_rollout( - PROCESS_SEGMENTS_TASK_ROLLOUT_OPTION + "spans.buffer.process-segments-task-rollout-rate" ): task_produce_future = process_segment_task.apply_async_with_future( args=[kafka_payload.value] From 3e1b2fa49891fedb8b7a651a3e9797d41b46ae51 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Mon, 6 Jul 2026 13:56:02 -0400 Subject: [PATCH 6/7] decide path for entire segment --- src/sentry/spans/consumers/process/flusher.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/sentry/spans/consumers/process/flusher.py b/src/sentry/spans/consumers/process/flusher.py index bc2bdd673276..043db8167a66 100644 --- a/src/sentry/spans/consumers/process/flusher.py +++ b/src/sentry/spans/consumers/process/flusher.py @@ -353,6 +353,11 @@ def produce(project_id: int, payload: KafkaPayload, dropped: int) -> None: }, ) + produce_to_process_segment_task = ( + produce_to_pipe is None + and in_random_rollout("spans.buffer.process-segments-task-rollout-rate") + ) + for message in flushed_segment.to_messages(): kafka_payload = KafkaPayload(None, orjson.dumps(message), []) metrics.timing( @@ -360,9 +365,7 @@ def produce(project_id: int, payload: KafkaPayload, dropped: int) -> None: len(kafka_payload.value), tags={"shard": shard_tag}, ) - if produce_to_pipe is None and in_random_rollout( - "spans.buffer.process-segments-task-rollout-rate" - ): + if produce_to_process_segment_task: task_produce_future = process_segment_task.apply_async_with_future( args=[kafka_payload.value] ) From 87047aca6fe80e89d22810488a41bbfff5d18993 Mon Sep 17 00:00:00 2001 From: Tony Le Date: Mon, 6 Jul 2026 13:59:39 -0400 Subject: [PATCH 7/7] fix: typing --- src/sentry/spans/consumers/process/flusher.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/sentry/spans/consumers/process/flusher.py b/src/sentry/spans/consumers/process/flusher.py index 043db8167a66..a2c6bbc5eb55 100644 --- a/src/sentry/spans/consumers/process/flusher.py +++ b/src/sentry/spans/consumers/process/flusher.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import Future from functools import partial +from typing import Any import orjson import sentry_sdk @@ -282,7 +283,7 @@ def main( logger.info("Flusher process started for shards %s", shard_tag) try: - producer_futures = [] + producer_futures: list[tuple[int, Any, int]] = [] if produce_to_pipe is not None: