ref(replays): Remove dead code and reorganize ingest-replay-recordings consumer (#95030)

Problem statement:

There was a variety of code paths which were no longer used. These were
dropped. Additionally some of the code was not well factored. Too many
concerns were put into one module and which reduced its ability to
integrate with multiple systems and also led to weird control flow for
cross module communication. This has been refactored and unit test and
integration test coverage has been added.

What changed:

- Message parsing, header parsing, and decompression were all moved to
the recording consumer module.
- The ingest use case focuses on ingesting the data received. Not the
format it was received in.
- 700 lines of unit test coverage added.
- A bunch of e2e test coverage was removed because it was slow and
didn't properly test the way the unit tests do.

Author: cmanallen

src/sentry/replays/consumers/recording.py

@@ -1,27 +1,41 @@
 import logging
+import zlib
 from collections.abc import Mapping
+from typing import cast
 
 import sentry_sdk
+import sentry_sdk.scope
 from arroyo.backends.kafka.consumer import KafkaPayload
 from arroyo.processing.strategies import RunTask, RunTaskInThreads
 from arroyo.processing.strategies.abstract import ProcessingStrategy, ProcessingStrategyFactory
 from arroyo.processing.strategies.commit import CommitOffsets
 from arroyo.types import Commit, FilteredPayload, Message, Partition
 from django.conf import settings
+from sentry_kafka_schemas.codecs import Codec, ValidationError
+from sentry_kafka_schemas.schema_types.ingest_replay_recordings_v1 import ReplayRecording
+from sentry_sdk import set_tag
 
+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 (
-    DropSilently,
-    ProcessedRecordingMessage,
+    DropEvent,
+    Event,
+    ProcessedEvent,
     commit_recording_message,
-    parse_recording_message,
-    process_recording_message,
+    process_recording_event,
     track_recording_metadata,
 )
+from sentry.utils import json
+
+RECORDINGS_CODEC: Codec[ReplayRecording] = get_topic_codec(Topic.INGEST_REPLAYS_RECORDINGS)
 
 logger = logging.getLogger(__name__)
 
 
+class DropSilently(Exception):
+    pass
+
+
 class ProcessReplayRecordingStrategyFactory(ProcessingStrategyFactory[KafkaPayload]):
     def __init__(
         self,
@@ -61,7 +75,10 @@ def create_with_partitions(
         )
 
 
-def process_message(message: Message[KafkaPayload]) -> ProcessedRecordingMessage | FilteredPayload:
+# Processing Task
+
+
+def process_message(message: Message[KafkaPayload]) -> ProcessedEvent | FilteredPayload:
     with sentry_sdk.start_transaction(
         name="replays.consumer.recording_buffered.process_message",
         op="replays.consumer.recording_buffered.process_message",
@@ -70,15 +87,87 @@ def process_message(message: Message[KafkaPayload]) -> ProcessedRecordingMessage
         },
     ):
         try:
-            return process_recording_message(parse_recording_message(message.payload.value))
+            recording_event = parse_recording_event(message.payload.value)
+            set_tag("org_id", recording_event["context"]["org_id"])
+            set_tag("project_id", recording_event["context"]["project_id"])
+            return process_recording_event(recording_event)
         except DropSilently:
             return FilteredPayload()
         except Exception:
             logger.exception("Failed to process replay recording message.")
             return FilteredPayload()
 
 
-def commit_message(message: Message[ProcessedRecordingMessage]) -> None:
+@sentry_sdk.trace
+def parse_recording_event(message: bytes) -> Event:
+    recording = parse_request_message(message)
+    segment_id, payload = parse_headers(cast(bytes, recording["payload"]), recording["replay_id"])
+    compressed, decompressed = decompress_segment(payload)
+
+    replay_event_json = recording.get("replay_event")
+    if replay_event_json:
+        replay_event = json.loads(cast(bytes, replay_event_json))
+    else:
+        replay_event = None
+
+    replay_video_raw = recording.get("replay_video")
+    if replay_video_raw is not None:
+        replay_video = cast(bytes, replay_video_raw)
+    else:
+        replay_video = None
+
+    return {
+        "context": {
+            "key_id": recording.get("key_id"),
+            "org_id": recording["org_id"],
+            "project_id": recording["project_id"],
+            "received": recording["received"],
+            "replay_id": recording["replay_id"],
+            "retention_days": recording["retention_days"],
+            "segment_id": segment_id,
+        },
+        "payload_compressed": compressed,
+        "payload": decompressed,
+        "replay_event": replay_event,
+        "replay_video": replay_video,
+    }
+
+
+@sentry_sdk.trace
+def parse_request_message(message: bytes) -> ReplayRecording:
+    try:
+        return RECORDINGS_CODEC.decode(message)
+    except ValidationError:
+        logger.exception("Could not decode recording message.")
+        raise DropSilently()
+
+
+@sentry_sdk.trace
+def decompress_segment(segment: bytes) -> tuple[bytes, bytes]:
+    try:
+        return (segment, zlib.decompress(segment))
+    except zlib.error:
+        if segment and segment[0] == ord("["):
+            return (zlib.compress(segment), segment)
+        else:
+            logger.exception("Invalid recording body.")
+            raise DropSilently()
+
+
+@sentry_sdk.trace
+def parse_headers(recording: bytes, replay_id: str) -> tuple[int, bytes]:
+    try:
+        recording_headers_json, recording_segment = recording.split(b"\n", 1)
+        return int(json.loads(recording_headers_json)["segment_id"]), recording_segment
+    except Exception:
+        logger.exception("Recording headers could not be extracted %s", replay_id)
+        raise DropSilently()
+
+
+# I/O Task
+
+
+def commit_message(message: Message[ProcessedEvent]) -> None:
     isolation_scope = sentry_sdk.get_isolation_scope().fork()
     with sentry_sdk.scope.use_isolation_scope(isolation_scope):
         with sentry_sdk.start_transaction(
@@ -96,7 +185,7 @@ def commit_message(message: Message[ProcessedRecordingMessage]) -> None:
                 return None
             except GCS_RETRYABLE_ERRORS:
                 raise
-            except DropSilently:
+            except DropEvent:
                 return None
             except Exception:
                 logger.exception("Failed to commit replay recording message.")