feat(replays): Add buffered consumer implementation

bin/mock-replay-recording

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+""".
+
+Helpful commands:
+
+    - Run the consumer.
+        - `sentry run consumer ingest-replay-recordings --consumer-group 0`
+    - Check if offsets are committed correctly.
+        - `docker exec -it kafka-kafka-1 kafka-consumer-groups --bootstrap-server localhost:9092 --describe --group 0`
+"""
+from sentry.runner import configure
+
+configure()
+import logging
+import os
+import time
+import uuid
+
+import click
+from arroyo import Topic as ArroyoTopic
+from arroyo.backends.kafka import KafkaPayload, KafkaProducer, build_kafka_configuration
+from sentry_kafka_schemas.codecs import Codec
+from sentry_kafka_schemas.schema_types.ingest_replay_recordings_v1 import ReplayRecording
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "sentry.conf.server")
+
+import django
+
+django.setup()
+
+from sentry.conf.types.kafka_definition import Topic, get_topic_codec
+from sentry.utils.kafka_config import get_kafka_producer_cluster_options, get_topic_definition
+
+logger = logging.getLogger(__name__)
+
+
+def get_producer() -> KafkaProducer:
+    cluster_name = get_topic_definition(Topic.INGEST_REPLAYS_RECORDINGS)["cluster"]
+    producer_config = get_kafka_producer_cluster_options(cluster_name)
+    return KafkaProducer(build_kafka_configuration(default_config=producer_config))
+
+
+RECORDING_CODEC: Codec[ReplayRecording] = get_topic_codec(Topic.INGEST_REPLAYS_RECORDINGS)
+
+
+@click.command()
+@click.option("--organization-id", type=int, required=True, help="Organization ID")
+@click.option("--project-id", type=int, required=True, help="Project ID")
+def main(organization_id: int, project_id: int) -> None:
+    """Produce a mock uptime result message to the INGEST_REPLAYS_RECORDINGS topic."""
+    message: ReplayRecording = {
+        "key_id": None,
+        "org_id": organization_id,
+        "payload": b'{"segment_id"',
+        "project_id": project_id,
+        "received": int(time.time()),
+        "replay_event": None,
+        "replay_id": uuid.uuid4().hex,
+        "replay_video": None,
+        "retention_days": 30,
+        "type": "replay_recording_not_chunked",
+        "version": 1,
+    }
+
+    producer = get_producer()
+    topic = get_topic_definition(Topic.INGEST_REPLAYS_RECORDINGS)["real_topic_name"]
+    payload = KafkaPayload(None, RECORDING_CODEC.encode(message), [])
+
+    producer.produce(ArroyoTopic(topic), payload)
+    producer.close()
+
+    logger.info("Successfully produced message to %s", topic)
+
+
+if __name__ == "__main__":
+    main()

src/sentry/consumers/__init__.py

@@ -84,24 +84,11 @@ def ingest_replay_recordings_options() -> list[click.Option]:
 
 def ingest_replay_recordings_buffered_options() -> list[click.Option]:
     """Return a list of ingest-replay-recordings-buffered options."""
-    options = [
-        click.Option(
-            ["--max-buffer-message-count", "max_buffer_message_count"],
-            type=int,
-            default=100,
-        ),
-        click.Option(
-            ["--max-buffer-size-in-bytes", "max_buffer_size_in_bytes"],
-            type=int,
-            default=2_500_000,
-        ),
-        click.Option(
-            ["--max-buffer-time-in-seconds", "max_buffer_time_in_seconds"],
-            type=int,
-            default=1,
-        ),
+    return [
+        click.Option(["--max-buffer-length", "max_buffer_length"], type=int, default=8),
+        click.Option(["--max-buffer-wait", "max_buffer_wait"], type=int, default=1),
+        click.Option(["--max-workers", "max_workers"], type=int, default=8),
     ]
-    return options
 
 
 def ingest_monitors_options() -> list[click.Option]:
@@ -269,8 +256,8 @@ def ingest_transactions_options() -> list[click.Option]:
     },
     "ingest-replay-recordings": {
         "topic": Topic.INGEST_REPLAYS_RECORDINGS,
-        "strategy_factory": "sentry.replays.consumers.recording.ProcessReplayRecordingStrategyFactory",
-        "click_options": ingest_replay_recordings_options(),
+        "strategy_factory": "sentry.replays.consumers.buffered.factory.PlatformStrategyFactory",
+        "click_options": ingest_replay_recordings_buffered_options(),
     },
     "ingest-replay-recordings-buffered": {
         "topic": Topic.INGEST_REPLAYS_RECORDINGS,

src/sentry/replays/consumers/buffered/consumer.py

@@ -0,0 +1,196 @@
+"""Session Replay recording consumer implementation.
+
+The consumer implementation follows a batching flush strategy. We accept messages, process them,
+buffer them, and when some threshold is reached we flush the buffer. The batch has finished work
+after the buffer is flushed so we commit with a None value.
+"""
+
+import contextlib
+import time
+from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
+from dataclasses import dataclass
+from typing import TypedDict
+
+import sentry_sdk
+
+from sentry.replays.consumers.buffered.platform import (
+    Cmd,
+    Commit,
+    Effect,
+    Join,
+    Nothing,
+    Poll,
+    RunTime,
+    Sub,
+    Task,
+)
+from sentry.replays.usecases.ingest import (
+    DropSilently,
+    ProcessedRecordingMessage,
+    commit_recording_message,
+    parse_recording_message,
+    process_recording_message,
+    track_recording_metadata,
+)
+
+# Types.
+
+
+class Flags(TypedDict):
+    max_buffer_length: int
+    max_buffer_wait: int
+    max_workers: int
+
+
+@dataclass
+class Model:
+    buffer: list[ProcessedRecordingMessage]
+    last_flushed_at: float
+    max_buffer_length: int
+    max_buffer_wait: int
+    max_workers: int
+
+
+@dataclass(frozen=True)
+class Append:
+    """Append the item to the buffer."""
+
+    item: ProcessedRecordingMessage
+
+
+class Committed:
+    """The platform committed offsets. Our buffer is now completely done."""
+
+
+class Flush:
+    """Our application hit the flush threshold and has been instructed to flush."""
+
+
+@dataclass(frozen=True)
+class Flushed:
+    """Our application successfully flushed."""
+
+    now: float
+
+
+class Skip:
+    """Skip the message."""
+
+
+@dataclass(frozen=True)
+class TryFlush:
+    """Instruct the application to flush the buffer if its time."""
+
+    now: float
+
+
+# A "Msg" is the union of all application messages our RunTime will accept.
+Msg = Append | Committed | Flush | Flushed | Skip | TryFlush
+
+
+# State machine functions.
+
+
+def init(flags: Flags) -> tuple[Model, Cmd[Msg, None]]:
+    return (
+        Model(
+            buffer=[],
+            last_flushed_at=time.time(),
+            max_buffer_wait=flags["max_buffer_wait"],
+            max_workers=flags["max_workers"],
+            max_buffer_length=flags["max_buffer_length"],
+        ),
+        Nothing(),
+    )
+
+
+@sentry_sdk.trace
+def process(_: Model, message: bytes) -> Msg:
+    try:
+        item = process_recording_message(parse_recording_message(message))
+        return Append(item=item)
+    except Exception:
+        return Skip()
+
+
+def update(model: Model, msg: Msg) -> tuple[Model, Cmd[Msg, None]]:
+    match msg:
+        case Append(item=item):
+            model.buffer.append(item)
+            return (model, Effect(fun=time.time, msg=lambda now: TryFlush(now=now)))
+        case Skip():
+            return (model, Effect(fun=time.time, msg=lambda now: TryFlush(now=now)))
+        case Committed():
+            return (model, Nothing())
+        case Flush():
+            return (model, Effect(fun=FlushBuffer(model), msg=lambda now: Flushed(now=now)))
+        case Flushed(now=now):
+            model.buffer = []
+            model.last_flushed_at = now
+            return (model, Commit(msg=Committed(), value=None))
+        case TryFlush(now=now):
+            return (model, Task(msg=Flush())) if can_flush(model, now) else (model, Nothing())
+
+
+def subscription(model: Model) -> list[Sub[Msg]]:
+    return [
+        Join(msg=Flush),
+        Poll(msg=lambda: TryFlush(now=time.time())),
+    ]
+
+
+# Helpers.
+
+
+def can_flush(model: Model, now: float) -> bool:
+    return (
+        len(model.buffer) >= model.max_buffer_length
+        or (now - model.max_buffer_wait) >= model.last_flushed_at
+    )
+
+
+@dataclass(frozen=True)
+class FlushBuffer:
+    model: Model
+
+    def __call__(self) -> float:
+        @sentry_sdk.trace
+        def flush_message(message: ProcessedRecordingMessage) -> None:
+            with contextlib.suppress(DropSilently):
+                commit_recording_message(message)
+
+        if len(self.model.buffer) == 0:
+            return time.time()
+
+        with ThreadPoolExecutor(max_workers=self.model.max_workers) as pool:
+            futures = [pool.submit(flush_message, message) for message in self.model.buffer]
+
+            # Tasks can fail. We check the done set for any failures. We will wait for all the
+            # futures to complete before running this step or eagerly run this step if any task
+            # errors.
+            done, _ = wait(futures, return_when=FIRST_EXCEPTION)
+            for future in done:
+                exc = future.exception()
+                # Raising preserves the existing behavior. We can address error handling in a
+                # follow up.
+                if exc is not None and not isinstance(exc, DropSilently):
+                    raise exc
+
+        # Recording metadata is not tracked in the threadpool. This is because this function will
+        # log. Logging will acquire a lock and make our threading less useful due to the speed of
+        # the I/O we do in this step.
+        for message in self.model.buffer:
+            track_recording_metadata(message)
+
+        return time.time()
+
+
+# Consumer.
+
+
+recording_consumer = RunTime(
+    init=init,
+    process=process,
+    subscription=subscription,
+    update=update,
+)

src/sentry/replays/consumers/buffered/factory.py

@@ -0,0 +1,34 @@
+"""Session Replay recording consumer strategy factory.
+
+This module exists solely to abstract the bootstraping process of the application and runtime in
+`sentry/consumers/__init__.py`.
+"""
+
+from collections.abc import Mapping
+
+from arroyo.backends.kafka.consumer import KafkaPayload
+from arroyo.processing.strategies import ProcessingStrategy, ProcessingStrategyFactory
+from arroyo.processing.strategies.commit import CommitOffsets
+from arroyo.types import Commit, Partition
+
+from sentry.replays.consumers.buffered.consumer import Flags, recording_consumer
+from sentry.replays.consumers.buffered.platform import PlatformStrategy
+
+
+class PlatformStrategyFactory(ProcessingStrategyFactory[KafkaPayload]):
+
+    def __init__(self, max_buffer_length: int, max_buffer_wait: int, max_workers: int) -> None:
+        self.flags: Flags = {
+            "max_buffer_length": max_buffer_length,
+            "max_buffer_wait": max_buffer_wait,
+            "max_workers": max_workers,
+        }
+
+    def create_with_partitions(
+        self,
+        commit: Commit,
+        partitions: Mapping[Partition, int],
+    ) -> ProcessingStrategy[KafkaPayload]:
+        return PlatformStrategy(
+            next_step=CommitOffsets(commit), flags=self.flags, runtime=recording_consumer
+        )