ref(producer): Use arroyo KafkaProducer for commit log

snuba/consumers/consumer.py

@@ -15,6 +15,7 @@
 
 import rapidjson
 import sentry_sdk
+from arroyo.backends.abstract import Producer
 from arroyo.backends.kafka import KafkaPayload
 from arroyo.backends.kafka.commit import CommitCodec
 from arroyo.commit import Commit as CommitLogCommit
@@ -23,7 +24,6 @@
 from confluent_kafka import KafkaError
 from confluent_kafka import Message as ConfluentMessage
 from confluent_kafka import Producer as ConfluentKafkaProducer
-from confluent_kafka import Producer as ConfluentProducer
 
 from snuba import environment
 from snuba.clickhouse.http import JSONRow, JSONRowEncoder, ValuesRowEncoder
@@ -49,7 +49,7 @@
 
 
 class CommitLogConfig(NamedTuple):
-    producer: ConfluentProducer
+    producer: Producer[KafkaPayload]
     topic: Topic
     group_id: str
 
@@ -270,12 +270,6 @@ def submit(self, message: Message[None | BytesInsertBatch | ReplacementBatch]) -
             message.value.timestamp,
         )
 
-    def __commit_message_delivery_callback(
-        self, error: KafkaError | None, message: ConfluentMessage
-    ) -> None:
-        if error is not None:
-            raise Exception(error.str())
-
     def close(self) -> None:
         self.__closed = True
 
@@ -285,6 +279,7 @@ def close(self) -> None:
             self.__replacement_batch_writer.close()
 
         if self.__commit_log_config is not None:
+            futures = []
             for partition, (
                 offset,
                 timestamp,
@@ -305,14 +300,14 @@ def close(self) -> None:
                         received_p99,
                     )
                 )
-                self.__commit_log_config.producer.produce(
-                    self.__commit_log_config.topic.name,
-                    key=payload.key,
-                    value=payload.value,
-                    headers=cast("list[tuple[str, str | bytes | None]]", payload.headers),
-                    on_delivery=self.__commit_message_delivery_callback,
+                futures.append(
+                    self.__commit_log_config.producer.produce(
+                        self.__commit_log_config.topic, payload
+                    )
                 )
-                self.__commit_log_config.producer.poll(0.0)
+            # raises on delivery failure, so offsets are not committed past a missing commit log entry
+            for future in futures:
+                future.result()
         self.__offsets_to_produce.clear()
         self.__received_timestamps.clear()
 
@@ -419,6 +414,7 @@ def close(self) -> None:
             step.close()
 
         if self.__commit_log_config is not None:
+            futures = []
             for partition, (offset, timestamp) in self.__offsets_to_produce.items():
                 payload = commit_codec.encode(
                     CommitLogCommit(
@@ -429,26 +425,17 @@ def close(self) -> None:
                         None,
                     )
                 )
-                self.__commit_log_config.producer.produce(
-                    self.__commit_log_config.topic.name,
-                    key=payload.key,
-                    value=payload.value,
-                    headers=cast("list[tuple[str, str | bytes | None]]", payload.headers),
-                    on_delivery=self.__commit_message_delivery_callback,
+                futures.append(
+                    self.__commit_log_config.producer.produce(
+                        self.__commit_log_config.topic, payload
+                    )
                 )
-                self.__commit_log_config.producer.poll(0.0)
-
-            self.__commit_log_config.producer.flush()
+            for future in futures:
+                future.result()
 
         self.__messages = {}
         self.__offsets_to_produce.clear()
 
-    def __commit_message_delivery_callback(
-        self, error: KafkaError | None, message: ConfluentMessage
-    ) -> None:
-        if error is not None:
-            raise Exception(error.str())
-
 
 class MultistorageKafkaPayload(NamedTuple):
     storage_keys: Sequence[StorageKey]

snuba/consumers/consumer_builder.py

@@ -126,9 +126,9 @@ def __init__(
             else None
         )
 
-        self.commit_log_producer: Producer | None
+        self.commit_log_producer: KafkaProducer | None
         if self.__consumer_config.commit_log_topic is not None:
-            self.commit_log_producer = Producer(
+            self.commit_log_producer = KafkaProducer(
                 build_kafka_configuration(self.__consumer_config.commit_log_topic.broker_config)
             )
         else:
@@ -317,7 +317,7 @@ def flush(self) -> None:
             self.replacements_producer.flush()
 
         if self.commit_log_producer:
-            self.commit_log_producer.flush()
+            self.commit_log_producer.close().result()
 
         if self.dlq_producer:
             self.dlq_producer.close()

tests/test_consumer.py

@@ -3,18 +3,21 @@
 import json
 import pickle
 from collections.abc import MutableSequence
+from concurrent.futures import Future
 from datetime import datetime
 from pathlib import Path
 from pickle import PickleBuffer
 from unittest.mock import Mock, call
 
 import pytest
 from arroyo.backends.kafka import KafkaPayload
+from arroyo.backends.kafka.commit import CommitCodec
 from arroyo.types import BrokerValue, Message, Partition, Topic
 
 from snuba.clusters.cluster import ClickhouseClientSettings
 from snuba.consumers.consumer import (
     BytesInsertBatch,
+    CommitLogConfig,
     InsertBatchWriter,
     LatencyRecorder,
     ProcessedMessageBatchWriter,
@@ -124,6 +127,39 @@ def get_number_of_insertion_metrics() -> int:
     assert health_check_file.exists()
 
 
+def test_commit_log_produced_on_close() -> None:
+    commit_log_topic = Topic("commit-log")
+    partition = Partition(Topic("events"), 0)
+
+    def close_with(future: Future[BrokerValue[KafkaPayload]]) -> Mock:
+        producer = Mock()
+        producer.produce.return_value = future
+        writer = ProcessedMessageBatchWriter(
+            insert_batch_writer=InsertBatchWriter(
+                Mock(), MetricsWrapper(TestingMetricsBackend(), "insertions")
+            ),
+            commit_log_config=CommitLogConfig(producer, commit_log_topic, "group"),
+        )
+        writer.submit(
+            Message(BrokerValue(BytesInsertBatch([b"{}"], None), partition, 5, datetime.now()))
+        )
+        writer.close()
+        return producer
+
+    delivered: Future[BrokerValue[KafkaPayload]] = Future()
+    delivered.set_result(Mock())
+    producer = close_with(delivered)
+    destination, payload = producer.produce.call_args.args
+    commit = CommitCodec().decode(payload)
+    assert destination == commit_log_topic
+    assert (commit.group, commit.partition, commit.offset) == ("group", partition, 6)
+
+    failed: Future[BrokerValue[KafkaPayload]] = Future()
+    failed.set_exception(Exception("delivery failed"))
+    with pytest.raises(Exception, match="delivery failed"):
+        close_with(failed)
+
+
 def test_json_row_batch_pickle_simple() -> None:
     batch = BytesInsertBatch([b"foo", b"bar", b"baz"], datetime(2021, 1, 1, 11, 0, 1))
     assert pickle.loads(pickle.dumps(batch)) == batch