diff --git a/snuba/state/__init__.py b/snuba/state/__init__.py index 953eff7007..d4b56b6350 100644 --- a/snuba/state/__init__.py +++ b/snuba/state/__init__.py @@ -2,14 +2,16 @@ import logging from collections.abc import Mapping +from concurrent.futures import Future from functools import partial from typing import ( Any, ) import simplejson as json -from confluent_kafka import KafkaError, Producer -from confluent_kafka import Message as KafkaMessage +from arroyo.backends.kafka import KafkaPayload, KafkaProducer +from arroyo.types import BrokerValue +from arroyo.types import Topic as ArroyoTopic from sentry_kafka_schemas.schema_types import snuba_queries_v1 from snuba import environment, settings @@ -20,7 +22,7 @@ metrics = MetricsWrapper(environment.metrics, "state") logger = logging.getLogger("snuba.state") -kfk: Producer | None = None +kfk: KafkaProducer | None = None ratelimit_prefix = "snuba-ratelimit:" @@ -32,10 +34,10 @@ rate_lookback_s = 60 -def _kafka_producer() -> Producer: +def _kafka_producer() -> KafkaProducer: global kfk if kfk is None: - kfk = Producer( + kfk = KafkaProducer( build_kafka_producer_configuration( topic=Topic.QUERYLOG, ) @@ -55,7 +57,8 @@ def safe_dumps_default(value: Any) -> Any: safe_dumps = partial(json.dumps, for_json=True, default=safe_dumps_default) -def _record_query_delivery_callback(error: KafkaError | None, message: KafkaMessage) -> None: +def _record_query_delivery_callback(future: Future[BrokerValue[KafkaPayload]]) -> None: + error = future.exception() metrics.increment( "record_query.delivery_callback", tags={"status": "success" if error is None else "failure"}, @@ -72,12 +75,12 @@ def record_query(query_metadata: snuba_queries_v1.Querylog) -> None: try: producer = _kafka_producer() data = safe_dumps(query_metadata) - producer.poll(0) # trigger queued delivery callbacks - producer.produce( - settings.KAFKA_TOPIC_MAP.get(Topic.QUERYLOG.value, Topic.QUERYLOG.value), - data.encode("utf-8"), - on_delivery=_record_query_delivery_callback, + future = producer.produce( + ArroyoTopic(settings.KAFKA_TOPIC_MAP.get(Topic.QUERYLOG.value, Topic.QUERYLOG.value)), + KafkaPayload(None, data.encode("utf-8"), []), ) + assert isinstance(future, Future) # use_simple_futures is off + future.add_done_callback(_record_query_delivery_callback) except Exception as ex: logger.exception("Could not record query due to error: %r", ex) @@ -85,5 +88,6 @@ def record_query(query_metadata: snuba_queries_v1.Querylog) -> None: def flush_producer() -> None: global kfk if kfk is not None: - messages_remaining = kfk.flush() - logger.debug(f"{messages_remaining} querylog messages pending delivery") + # close() waits for in-flight messages to be delivered + kfk.close().result() + kfk = None diff --git a/snuba/web/bulk_delete_query.py b/snuba/web/bulk_delete_query.py index 17d4b656dd..523dcbd4a0 100644 --- a/snuba/web/bulk_delete_query.py +++ b/snuba/web/bulk_delete_query.py @@ -1,12 +1,12 @@ import logging -import time from collections.abc import Mapping, MutableMapping, Sequence -from threading import Thread +from concurrent.futures import Future from typing import Any, TypedDict import rapidjson -from confluent_kafka import KafkaError, Producer -from confluent_kafka import Message as KafkaMessage +from arroyo.backends.kafka import KafkaPayload, KafkaProducer +from arroyo.types import BrokerValue +from arroyo.types import Topic as ArroyoTopic from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType from snuba import environment, settings @@ -61,7 +61,7 @@ class DeleteQueryMessage(TypedDict, total=False): attribute_conditions_item_type: TraceItemType.ValueType | None -PRODUCER_MAP: MutableMapping[str, Producer] = {} +PRODUCER_MAP: MutableMapping[str, KafkaProducer] = {} STORAGE_TOPIC: Mapping[str, Topic] = { StorageKey.SEARCH_ISSUES.value: Topic.LW_DELETIONS_GENERIC_EVENTS, # TODO: remove EAP_ITEMS once every region's consumer runs --storage eap_items_2 @@ -74,10 +74,11 @@ class InvalidStorageTopic(Exception): pass -def _get_kafka_producer(topic: Topic) -> Producer: +def _get_kafka_producer(topic: Topic) -> KafkaProducer: + # KafkaProducer polls in its own thread, so delivery callbacks fire without manual flushing producer = PRODUCER_MAP.get(topic.value) if not producer: - producer = Producer( + producer = KafkaProducer( build_kafka_producer_configuration( topic=topic, ) @@ -86,25 +87,8 @@ def _get_kafka_producer(topic: Topic) -> Producer: return producer -def flush_producers() -> None: - """ - It's not guaranteed that there will be a steady stream of - DELETE requests so we can call producer.flush() for any active - producers to make sure the delivery callbacks are called. - """ - - def _flush_producers() -> None: - while True: - for storage, producer in PRODUCER_MAP.items(): - messages_remaining = producer.flush(5.0) - if messages_remaining: - logger.debug(f"{messages_remaining} {storage} messages pending delivery") - time.sleep(1) - - Thread(target=_flush_producers, name="flush_producers", daemon=True).start() - - -def _delete_query_delivery_callback(error: KafkaError | None, message: KafkaMessage) -> None: +def _delete_query_delivery_callback(future: Future[BrokerValue[KafkaPayload]]) -> None: + error = future.exception() metrics.increment( "delete_query.delivery_callback", tags={"status": "failure" if error else "success"}, @@ -114,15 +98,7 @@ def _delete_query_delivery_callback(error: KafkaError | None, message: KafkaMess logger.warning("Could not produce delete query due to error: %r", error) -FLUSH_PRODUCERS_THREAD_STARTED = False - - def produce_delete_query(delete_query: DeleteQueryMessage) -> None: - global FLUSH_PRODUCERS_THREAD_STARTED - if not FLUSH_PRODUCERS_THREAD_STARTED: - FLUSH_PRODUCERS_THREAD_STARTED = True - flush_producers() - storage_name = delete_query["storage_name"] topic = STORAGE_TOPIC.get(storage_name) if not topic: @@ -130,12 +106,12 @@ def produce_delete_query(delete_query: DeleteQueryMessage) -> None: try: producer = _get_kafka_producer(topic) data = rapidjson.dumps(delete_query).encode("utf-8") - producer.poll(0) # trigger queued delivery callbacks - producer.produce( - settings.KAFKA_TOPIC_MAP.get(topic.value, topic.value), - data, - on_delivery=_delete_query_delivery_callback, + future = producer.produce( + ArroyoTopic(settings.KAFKA_TOPIC_MAP.get(topic.value, topic.value)), + KafkaPayload(None, data, []), ) + assert isinstance(future, Future) # use_simple_futures is off + future.add_done_callback(_delete_query_delivery_callback) except Exception as ex: logger.exception("Could not produce delete query due to error: %r", ex)