diff --git a/snuba/consumers/consumer.py b/snuba/consumers/consumer.py index f98241dec4..a6ccef7610 100644 --- a/snuba/consumers/consumer.py +++ b/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] diff --git a/snuba/consumers/consumer_builder.py b/snuba/consumers/consumer_builder.py index e583fc4f5a..6512e43353 100644 --- a/snuba/consumers/consumer_builder.py +++ b/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() diff --git a/tests/test_consumer.py b/tests/test_consumer.py index c66eb6c92e..f157efd7e3 100644 --- a/tests/test_consumer.py +++ b/tests/test_consumer.py @@ -3,6 +3,7 @@ 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 @@ -10,11 +11,13 @@ 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