feat(consumers): add ProcessingStrategy for profiling join (#102140)

This PR adds a StrategyFactory which adds a strategy to a consumer that
starts/stops profiling during calls to `join()`.

The goal of profiling `join()` is to troubleshoot why some of our
consumers are hanging when `join()` is called during partition
revocation.

Author: bmckerry

src/sentry/consumers/__init__.py

@@ -23,6 +23,7 @@
     validate_consumer_definition,
 )
 from sentry.consumers.dlq import DlqStaleMessagesStrategyFactoryWrapper, maybe_build_dlq_producer
+from sentry.consumers.profiler import JoinProfiler
 from sentry.consumers.validate_schema import ValidateSchema
 from sentry.eventstream.types import EventStreamEventType
 from sentry.ingest.types import ConsumerType
@@ -466,6 +467,7 @@ def get_stream_processor(
     kafka_slice_id: int | None = None,
     shutdown_strategy_before_consumer: bool = False,
     add_global_tags: bool = False,
+    profile_consumer_join: bool = False,
 ) -> StreamProcessor:
     from sentry.utils import kafka_config
 
@@ -594,6 +596,9 @@ def build_consumer_config(group_id: str):
             healthcheck_file_path, strategy_factory
         )
 
+    if profile_consumer_join:
+        strategy_factory = JoinProfilerStrategyFactoryWrapper(strategy_factory)
+
     if enable_dlq and consumer_definition.get("dlq_topic"):
         dlq_topic = consumer_definition["dlq_topic"]
     else:
@@ -674,3 +679,12 @@ def __init__(self, healthcheck_file_path: str, inner: ProcessingStrategyFactory)
     def create_with_partitions(self, commit, partitions):
         rv = self.inner.create_with_partitions(commit, partitions)
         return Healthcheck(self.healthcheck_file_path, rv)
+
+
+class JoinProfilerStrategyFactoryWrapper(ProcessingStrategyFactory):
+    def __init__(self, inner: ProcessingStrategyFactory):
+        self.inner = inner
+
+    def create_with_partitions(self, commit, partitions):
+        rv = self.inner.create_with_partitions(commit, partitions)
+        return JoinProfiler(rv)

src/sentry/consumers/profiler.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import sentry_sdk
+from arroyo.processing.strategies.abstract import ProcessingStrategy
+from arroyo.types import Message, TStrategyPayload
+
+
+class JoinProfiler(ProcessingStrategy[TStrategyPayload]):
+    """
+    Strategy which passes through all ProcessingStrategy method calls,
+    but runs Sentry's continuous profiler for `join()` calls.
+
+    This startegy is being used to troubleshoot our consumers hanging during `join()` occasionally.
+    """
+
+    def __init__(self, next_step: ProcessingStrategy[TStrategyPayload]) -> None:
+        self.__next_step = next_step
+
+    def join(self, timeout: float | None = None):
+        with sentry_sdk.start_transaction(
+            op="consumer_join", name="consumer.join", custom_sampling_context={"sample_rate": 1.0}
+        ):
+            self.__next_step.join(timeout)
+
+    def submit(self, message: Message[TStrategyPayload]) -> None:
+        self.__next_step.submit(message)
+
+    def poll(self) -> None:
+        self.__next_step.poll()
+
+    def close(self) -> None:
+        self.__next_step.close()
+
+    def terminate(self) -> None:
+        self.__next_step.terminate()

src/sentry/options/defaults.py

@@ -3633,7 +3633,6 @@
 )
 
 # Sets the sample rate for profiles collected via the JoinProfiler arroyo strategy
-# TODO: add the JoinProfiler arroyo strategy
 register(
     "consumer.join.profiling.rate",
     type=Float,

src/sentry/runner/commands/run.py

@@ -544,6 +544,12 @@ def taskbroker_send_tasks(
     default=False,
     help="A potential workaround for Broker Handle Destroyed during shutdown (see arroyo option).",
 )
+@click.option(
+    "--profile-consumer-join",
+    is_flag=True,
+    default=False,
+    help="Adds a ProcessingStrategy to the start of a consumer that records a transaction of the consumer's join() method.",
+)
 @configuration
 def basic_consumer(
     consumer_name: str,
@@ -633,6 +639,7 @@ def dev_consumer(consumer_names: tuple[str, ...]) -> None:
             stale_threshold_sec=None,
             healthcheck_file_path=None,
             enforce_schema=True,
+            profile_consumer_join=False,
         )
         for consumer_name in consumer_names
     ]

src/sentry/utils/sdk.py

@@ -204,6 +204,7 @@ def traces_sampler(sampling_context):
 
 def profiles_sampler(sampling_context):
     PROFILES_SAMPLING_RATE = {
+        "consumer.join": options.get("consumer.join.profiling.rate"),
         "spans.process.process_message": options.get("spans.process-spans.profiling.rate"),
     }
     if "transaction_context" in sampling_context:

tests/sentry/consumers/test_profiler.py

@@ -0,0 +1,15 @@
+from unittest.mock import Mock
+
+from sentry.consumers import JoinProfilerStrategyFactoryWrapper
+
+
+def test_join_profiler() -> None:
+    inner_factory_mock = Mock()
+    inner_strategy_mock = Mock()
+    inner_factory_mock.create_with_partitions = Mock(return_value=inner_strategy_mock)
+
+    factory = JoinProfilerStrategyFactoryWrapper(inner_factory_mock)
+    strategy = factory.create_with_partitions(Mock(), Mock())
+
+    strategy.join()
+    inner_strategy_mock.join.assert_called_once()