fixes(REPLAY-592): Add continuous profiling to replay recording consumer (#97496)

fixes REPLAY-592

Add continuous profiling to process_message and commit_message in replay
recording consumer. Add an option
`replay.consumer.recording.profiling.enabled` to defaults.py, disabling
profiling by default. Add new tests for profiling.

Author: srest2021

src/sentry/options/defaults.py

@@ -524,6 +524,13 @@
     default=0.0,
     flags=FLAG_ALLOW_EMPTY | FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE,
 )
+# Enables profiling for replay recording ingestion.
+register(
+    "replay.consumer.recording.profiling.enabled",
+    type=Bool,
+    default=False,
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
 
 # User Feedback Options
 register(

src/sentry/replays/consumers/recording.py

@@ -15,6 +15,7 @@
 from sentry_kafka_schemas.schema_types.ingest_replay_recordings_v1 import ReplayRecording
 from sentry_sdk import set_tag
 
+from sentry import options
 from sentry.conf.types.kafka_definition import Topic, get_topic_codec
 from sentry.filestore.gcs import GCS_RETRYABLE_ERRORS
 from sentry.replays.usecases.ingest import (
@@ -65,9 +66,9 @@ def create_with_partitions(
         partitions: Mapping[Partition, int],
     ) -> ProcessingStrategy[KafkaPayload]:
         return RunTask(
-            function=process_message,
+            function=process_message_with_options,
             next_step=RunTaskInThreads(
-                processing_function=commit_message,
+                processing_function=commit_message_with_options,
                 concurrency=self.num_threads,
                 max_pending_futures=self.max_pending_futures,
                 next_step=CommitOffsets(commit),
@@ -78,14 +79,28 @@ def create_with_partitions(
 # Processing Task
 
 
-def process_message(message: Message[KafkaPayload]) -> ProcessedEvent | FilteredPayload:
+def process_message_with_options(
+    message: Message[KafkaPayload],
+) -> ProcessedEvent | FilteredPayload:
+    profiling_enabled = options.get(
+        "replay.consumer.recording.profiling.enabled",
+    )
+    return process_message(message, profiling_enabled=profiling_enabled)
+
+
+def process_message(
+    message: Message[KafkaPayload], profiling_enabled: bool = False
+) -> ProcessedEvent | FilteredPayload:
     with sentry_sdk.start_transaction(
         name="replays.consumer.recording_buffered.process_message",
         op="replays.consumer.recording_buffered.process_message",
         custom_sampling_context={
             "sample_rate": getattr(settings, "SENTRY_REPLAY_RECORDINGS_CONSUMER_APM_SAMPLING", 0)
         },
     ):
+        if profiling_enabled:
+            sentry_sdk.profiler.start_profiler()
+
         try:
             recording_event = parse_recording_event(message.payload.value)
             set_tag("org_id", recording_event["context"]["org_id"])
@@ -96,6 +111,9 @@ def process_message(message: Message[KafkaPayload]) -> ProcessedEvent | Filtered
         except Exception:
             logger.exception("Failed to process replay recording message.")
             return FilteredPayload()
+        finally:
+            if profiling_enabled:
+                sentry_sdk.profiler.stop_profiler()
 
 
 @sentry_sdk.trace
@@ -170,7 +188,14 @@ def parse_headers(recording: bytes, replay_id: str) -> tuple[int, bytes]:
 # I/O Task
 
 
-def commit_message(message: Message[ProcessedEvent]) -> None:
+def commit_message_with_options(message: Message[ProcessedEvent]) -> None:
+    profiling_enabled = options.get(
+        "replay.consumer.recording.profiling.enabled",
+    )
+    return commit_message(message, profiling_enabled=profiling_enabled)
+
+
+def commit_message(message: Message[ProcessedEvent], profiling_enabled: bool = False) -> None:
     isolation_scope = sentry_sdk.get_isolation_scope().fork()
     with sentry_sdk.scope.use_isolation_scope(isolation_scope):
         with sentry_sdk.start_transaction(
@@ -182,6 +207,9 @@ def commit_message(message: Message[ProcessedEvent]) -> None:
                 )
             },
         ):
+            if profiling_enabled:
+                sentry_sdk.profiler.start_profiler()
+
             try:
                 commit_recording_message(message.payload)
                 track_recording_metadata(message.payload)
@@ -193,3 +221,6 @@ def commit_message(message: Message[ProcessedEvent]) -> None:
             except Exception:
                 logger.exception("Failed to commit replay recording message.")
                 return None
+            finally:
+                if profiling_enabled:
+                    sentry_sdk.profiler.stop_profiler()

tests/sentry/replays/integration/consumers/test_recording.py

@@ -41,7 +41,11 @@ def submit(consumer, message):
     consumer.terminate()
 
 
-def test_recording_consumer(consumer) -> None:
+@mock.patch("sentry.replays.consumers.recording.options.get")
+def test_recording_consumer(mock_options_get, consumer) -> None:
+    # disable profiling
+    mock_options_get.return_value = False
+
     headers = json.dumps({"segment_id": 42}).encode()
     recording_payload = headers + b"\n" + zlib.compress(b"")
 
@@ -66,7 +70,11 @@ def test_recording_consumer(consumer) -> None:
         assert commit.called
 
 
-def test_recording_consumer_invalid_message(consumer) -> None:
+@mock.patch("sentry.replays.consumers.recording.options.get")
+def test_recording_consumer_invalid_message(mock_options_get, consumer) -> None:
+    # disable profiling
+    mock_options_get.return_value = False
+
     with mock.patch("sentry.replays.consumers.recording.commit_recording_message") as commit:
         submit(consumer, {})
 

tests/sentry/replays/unit/consumers/test_recording.py

@@ -1,4 +1,5 @@
 import zlib
+from unittest.mock import patch
 
 import msgpack
 import pytest
@@ -7,11 +8,14 @@
 
 from sentry.replays.consumers.recording import (
     DropSilently,
+    commit_message,
+    commit_message_with_options,
     decompress_segment,
     parse_headers,
     parse_recording_event,
     parse_request_message,
     process_message,
+    process_message_with_options,
 )
 from sentry.replays.usecases.ingest import ProcessedEvent
 from sentry.replays.usecases.ingest.event_parser import ParsedEventMeta
@@ -520,3 +524,148 @@ def test_process_message_no_headers() -> None:
 
 def make_kafka_message(message) -> Message[KafkaPayload]:
     return Message(Value(KafkaPayload(key=None, value=msgpack.packb(message), headers=[]), {}))
+
+
+def make_processed_event_message(processed_event: ProcessedEvent) -> Message[ProcessedEvent]:
+    return Message(Value(processed_event, {}))
+
+
+def make_valid_message() -> dict:
+    original_payload = b'[{"type": "test", "data": "some event data"}]'
+    compressed_payload = zlib.compress(original_payload)
+    segment_id = 42
+    headers = json.dumps({"segment_id": segment_id}).encode()
+    recording_payload = headers + b"\n" + compressed_payload
+
+    return {
+        "type": "replay_recording_not_chunked",
+        "org_id": 3,
+        "project_id": 4,
+        "replay_id": "1",
+        "received": 2,
+        "retention_days": 30,
+        "payload": recording_payload,
+        "key_id": 1,
+        "replay_event": b"{}",
+        "replay_video": b"",
+        "version": 0,
+    }
+
+
+def make_valid_processed_event() -> ProcessedEvent:
+    original_payload = b'[{"type": "test", "data": "some event data"}]'
+    compressed_payload = zlib.compress(original_payload)
+
+    return ProcessedEvent(
+        actions_event=ParsedEventMeta([], [], [], [], [], []),
+        context={
+            "key_id": 1,
+            "org_id": 3,
+            "project_id": 4,
+            "received": 2,
+            "replay_id": "1",
+            "retention_days": 30,
+            "segment_id": 42,
+        },
+        filedata=compressed_payload,
+        filename="30/4/1/42",
+        recording_size_uncompressed=len(original_payload),
+        recording_size=len(compressed_payload),
+        replay_event={},
+        trace_items=[],
+        video_size=None,
+    )
+
+
+@patch("sentry.replays.consumers.recording.sentry_sdk.profiler")
+@pytest.mark.parametrize("profiling_enabled", [True, False])
+def test_process_message_profiling(mock_profiler, profiling_enabled) -> None:
+    """Test that profiling is started and stopped when enabled, and not when disabled."""
+    message = make_valid_message()
+    kafka_message = make_kafka_message(message)
+    result = process_message(kafka_message, profiling_enabled=profiling_enabled)
+
+    assert mock_profiler.start_profiler.call_count == (1 if profiling_enabled else 0)
+    assert mock_profiler.stop_profiler.call_count == (1 if profiling_enabled else 0)
+
+    assert isinstance(result, ProcessedEvent)
+
+
+@patch("sentry.replays.consumers.recording.sentry_sdk.profiler")
+@pytest.mark.parametrize("profiling_enabled", [True, False])
+def test_commit_message_profiling(mock_profiler, profiling_enabled) -> None:
+    """Test that profiling is started and stopped when enabled, and not when disabled."""
+    processed_event = make_valid_processed_event()
+    commit_message(
+        make_processed_event_message(processed_event), profiling_enabled=profiling_enabled
+    )
+
+    assert mock_profiler.start_profiler.call_count == (1 if profiling_enabled else 0)
+    assert mock_profiler.stop_profiler.call_count == (1 if profiling_enabled else 0)
+
+
+@patch("sentry.replays.consumers.recording.sentry_sdk.profiler")
+@patch("sentry.replays.consumers.recording.parse_recording_event")
+@pytest.mark.parametrize("profiling_enabled", [True, False])
+def test_process_message_profiling_on_error(
+    mock_parse_recording_event, mock_profiler, profiling_enabled
+) -> None:
+    """Test that profiling is started and stopped when enabled, and not when disabled, even on error."""
+    mock_parse_recording_event.side_effect = Exception("test error")
+
+    message = make_valid_message()
+    kafka_message = make_kafka_message(message)
+    result = process_message(kafka_message, profiling_enabled=profiling_enabled)
+
+    assert mock_profiler.start_profiler.call_count == (1 if profiling_enabled else 0)
+    assert mock_profiler.stop_profiler.call_count == (1 if profiling_enabled else 0)
+    assert result == FilteredPayload()
+
+
+@patch("sentry.replays.consumers.recording.sentry_sdk.profiler")
+@patch("sentry.replays.consumers.recording.commit_recording_message")
+@pytest.mark.parametrize("profiling_enabled", [True, False])
+def test_commit_message_profiling_on_error(
+    mock_commit_recording_message, mock_profiler, profiling_enabled
+) -> None:
+    """Test that profiling is started and stopped when enabled, and not when disabled, even on error."""
+    mock_commit_recording_message.side_effect = Exception("test error")
+
+    processed_event = make_valid_processed_event()
+    commit_message(
+        make_processed_event_message(processed_event), profiling_enabled=profiling_enabled
+    )
+
+    assert mock_profiler.start_profiler.call_count == (1 if profiling_enabled else 0)
+    assert mock_profiler.stop_profiler.call_count == (1 if profiling_enabled else 0)
+
+
+@patch("sentry.replays.consumers.recording.options.get")
+@patch("sentry.replays.consumers.recording.commit_message")
+@pytest.mark.parametrize("profiling_enabled", [True, False])
+def test_commit_message_with_options(mock_commit_message, mock_options, profiling_enabled) -> None:
+    """Test that commit_message_with_options calls commit_message with the correct profiling_enabled value."""
+    mock_options.return_value = profiling_enabled
+
+    processed_event = make_valid_processed_event()
+    commit_message_with_options(make_processed_event_message(processed_event))
+
+    assert mock_commit_message.call_count == 1
+    assert mock_commit_message.call_args[1]["profiling_enabled"] == profiling_enabled
+
+
+@patch("sentry.replays.consumers.recording.options.get")
+@patch("sentry.replays.consumers.recording.process_message")
+@pytest.mark.parametrize("profiling_enabled", [True, False])
+def test_process_message_with_options(
+    mock_process_message, mock_options, profiling_enabled
+) -> None:
+    """Test that process_message_with_options calls process_message with the correct profiling_enabled value."""
+    mock_options.return_value = profiling_enabled
+
+    message = make_valid_message()
+    kafka_message = make_kafka_message(message)
+    process_message_with_options(kafka_message)
+
+    assert mock_process_message.call_count == 1
+    assert mock_process_message.call_args[1]["profiling_enabled"] == profiling_enabled