feat(producer): Replace confluent_kafka.Producer with ConfluentProducer (#102432)

Bump `sentry-arroyo` to `2.31.2` and slowly rolling out `ConfluentProducer`
to replace `confluent_kafka.Producer` in order to record produce metrics.

Refs STREAM-549

Author: lvthanh03

pyproject.toml

@@ -78,7 +78,7 @@ dependencies = [
     "rfc3339-validator>=0.1.2",
     "rfc3986-validator>=0.1.1",
     # [end] jsonschema format validators
-    "sentry-arroyo>=2.30.0",
+    "sentry-arroyo>=2.31.2",
     "sentry-forked-email-reply-parser>=0.5.12.post1",
     "sentry-kafka-schemas>=2.1.13",
     "sentry-ophio>=1.1.3",

src/sentry/eventstream/kafka/backend.py

@@ -20,6 +20,7 @@
 from sentry.eventstream.types import EventStreamEventType
 from sentry.killswitches import killswitch_matches_context
 from sentry.utils import json
+from sentry.utils.confluent_producer import get_confluent_producer
 from sentry.utils.kafka_config import get_kafka_producer_cluster_options, get_topic_definition
 
 EAP_ITEMS_CODEC: Codec[TraceItem] = get_topic_codec(Topic.SNUBA_ITEMS)
@@ -48,7 +49,7 @@ def get_producer(self, topic: Topic) -> Producer:
             cluster_options = get_kafka_producer_cluster_options(cluster_name)
             cluster_options["client.id"] = "sentry.eventstream.kafka"
             # XXX(markus): We should use `sentry.utils.arroyo_producer.get_arroyo_producer`.
-            self.__producers[topic] = Producer(
+            self.__producers[topic] = get_confluent_producer(
                 build_kafka_producer_configuration(default_config=cluster_options)
             )
 

src/sentry/options/defaults.py

@@ -629,6 +629,15 @@
     flags=FLAG_AUTOMATOR_MODIFIABLE,
 )
 
+# Rollout configuration for Arroyo ConfluentProducer
+# Controls the rollout of individual Kafka producers by name
+register(
+    "arroyo.producer.confluent-producer-rollout",
+    type=Dict,
+    default={},
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
+
 # Analytics
 register("analytics.backend", default="noop", flags=FLAG_NOSTORE)
 register("analytics.options", default={}, flags=FLAG_NOSTORE)

src/sentry/sentry_metrics/consumers/indexer/multiprocess.py

@@ -10,10 +10,10 @@
 from arroyo.processing.strategies import ProcessingStrategy as ProcessingStep
 from arroyo.processing.strategies.commit import CommitOffsets
 from arroyo.types import Commit, FilteredPayload, Message, Partition, Value
-from confluent_kafka import Producer
 
 from sentry.conf.types.kafka_definition import Topic
 from sentry.utils import kafka_config, metrics
+from sentry.utils.confluent_producer import get_confluent_producer
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +28,7 @@ def __init__(
         snuba_metrics = kafka_config.get_topic_definition(output_topic)
         producer_config = kafka_config.get_kafka_producer_cluster_options(snuba_metrics["cluster"])
         producer_config["client.id"] = "sentry.sentry_metrics.multiprocess"
-        self.__producer = Producer(producer_config)
+        self.__producer = get_confluent_producer(producer_config)
         self.__producer_topic = snuba_metrics["real_topic_name"]
 
         self.__commit = CommitOffsets(commit_function)

src/sentry/sentry_metrics/consumers/indexer/slicing_router.py

@@ -20,6 +20,7 @@
     RoutingPayload,
 )
 from sentry.utils import kafka_config
+from sentry.utils.confluent_producer import get_confluent_producer
 
 
 class SlicingConfigurationException(Exception):
@@ -112,8 +113,9 @@ def __init__(
             producer_config["client.id"] = (
                 f"sentry.sentry_metrics.slicing_router.{current_sliceable}.{current_slice_id}"
             )
+            producer = get_confluent_producer(producer_config)
             self.__slice_to_producer[current_slice_id] = MessageRoute(
-                producer=Producer(producer_config),
+                producer=producer,
                 topic=Topic(configuration["topic"]),
             )
         # All logical partitions should be routed to a slice ID that's present in the slice

src/sentry/utils/confluent_producer.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Any
+
+from arroyo.backends.kafka import ConfluentProducer
+from confluent_kafka import Producer
+
+from sentry import options
+
+
+def get_confluent_producer(
+    configuration: dict[str, Any],
+) -> Producer:
+    """
+    Get a confluent_kafka Producer for a given configuration.
+
+    Args:
+        configuration: The configuration for the confluent_kafka Producer
+
+    Returns:
+        confluent_kafka Producer
+    """
+    rollout_config = options.get("arroyo.producer.confluent-producer-rollout", {})
+    name = configuration.get("client.id")
+
+    # f"sentry.sentry_metrics.slicing_router.{current_sliceable}.{current_slice_id}"
+    if name is not None and isinstance(name, str):
+        if name.startswith("sentry.sentry_metrics.slicing_router"):
+            name = "sentry.sentry_metrics.slicing_router"
+
+    if rollout_config.get(name, False):
+        return ConfluentProducer(configuration)
+    return Producer(configuration)

src/sentry/utils/pubsub.py

@@ -1,14 +1,18 @@
 from typing import Any
 
 from arroyo.backends.kafka import build_kafka_producer_configuration
-from confluent_kafka import Producer
+
+from sentry.utils.confluent_producer import get_confluent_producer
 
 
 class KafkaPublisher:
     # XXX(markus): Deprecated. Please use `sentry.utils.arroyo_producer.get_arroyo_producer`.
     def __init__(self, connection: dict[str, Any], asynchronous: bool = True) -> None:
-        self.producer = Producer(
-            build_kafka_producer_configuration(default_config=connection or {})
+        connection = connection or {}
+        if "client.id" not in connection:
+            connection["client.id"] = "sentry.utils.pubsub"
+        self.producer = get_confluent_producer(
+            build_kafka_producer_configuration(default_config=connection)
         )
         self.asynchronous = asynchronous
 

tests/sentry/sentry_metrics/consumers/test_slicing_router.py

@@ -20,6 +20,7 @@
     _validate_slicing_config,
     _validate_slicing_consumer_config,
 )
+from sentry.testutils.helpers.options import override_options
 
 KAFKA_SNUBA_GENERIC_METRICS = "snuba-generic-metrics"
 
@@ -78,6 +79,7 @@ def setup_slicing() -> Generator[None]:
 
 
 @pytest.mark.parametrize("org_id", [1, 127, 128, 256, 257])
+@override_options({"arroyo.producer.confluent-producer-rollout": {}})
 def test_with_slicing(metrics_message, setup_slicing) -> None:
     """
     With partitioning settings, the SlicingRouter should route to the correct topic
@@ -94,6 +96,7 @@ def test_with_slicing(metrics_message, setup_slicing) -> None:
         assert False, "unexpected org_id"
 
 
+@override_options({"arroyo.producer.confluent-producer-rollout": {}})
 def test_with_no_org_in_routing_header(setup_slicing) -> None:
     """
     With partitioning settings, the SlicingRouter should route to the correct topic

tests/sentry/utils/test_confluent_producer.py

@@ -0,0 +1,45 @@
+from arroyo.backends.kafka import ConfluentProducer
+from confluent_kafka import Producer
+
+from sentry.testutils.helpers.options import override_options
+from sentry.utils.confluent_producer import get_confluent_producer
+
+
+@override_options({"arroyo.producer.confluent-producer-rollout": {"test_producer": True}})
+def test_get_confluent_producer() -> None:
+    configuration = {
+        "bootstrap.servers": "localhost:9092",
+        "client.id": "test_producer",
+    }
+
+    producer = get_confluent_producer(configuration)
+    assert producer is not None
+
+    assert isinstance(producer, Producer)
+    assert isinstance(producer, ConfluentProducer)
+
+    for attrs in Producer.__dict__.keys():
+        assert hasattr(producer, attrs)
+
+
+@override_options({"arroyo.producer.confluent-producer-rollout": {"test_producer": False}})
+def test_get_confluent_producer_not_rolled_out() -> None:
+    configuration = {
+        "bootstrap.servers": "localhost:9092",
+        "client.id": "test_producer",
+    }
+
+    producer = get_confluent_producer(configuration)
+    assert isinstance(producer, Producer)
+    assert not isinstance(producer, ConfluentProducer)
+
+
+@override_options({"arroyo.producer.confluent-producer-rollout": {}})
+def test_get_confluent_producer_no_client_id() -> None:
+    configuration = {
+        "bootstrap.servers": "localhost:9092",
+    }
+
+    producer = get_confluent_producer(configuration)
+    assert isinstance(producer, Producer)
+    assert not isinstance(producer, ConfluentProducer)