perf: Arroyo 2.15.0 and reusable multiprocessing pools

requirements.txt

@@ -22,7 +22,7 @@ pytest-watch==4.2.0
 python-dateutil==2.8.2
 python-rapidjson==1.8
 redis==4.3.4
-sentry-arroyo==2.14.25
+sentry-arroyo==2.15.0
 sentry-kafka-schemas==0.1.37
 sentry-redis-tools==0.1.7
 sentry-relay==0.8.27

snuba/consumers/consumer.py

@@ -57,6 +57,7 @@
 from snuba.clickhouse.http import JSONRow, JSONRowEncoder, ValuesRowEncoder
 from snuba.consumers.schemas import _NOOP_CODEC, get_json_codec
 from snuba.consumers.types import KafkaMessageMetadata
+from snuba.consumers.utils import get_reusable_multiprocessing_pool
 from snuba.datasets.storage import WritableTableStorage
 from snuba.datasets.storages.factory import get_writable_storage
 from snuba.datasets.storages.storage_key import StorageKey
@@ -389,7 +390,6 @@ def build_batch_writer(
     commit_log_config: Optional[CommitLogConfig] = None,
     slice_id: Optional[int] = None,
 ) -> Callable[[], ProcessedMessageBatchWriter]:
-
     assert not (replacements_producer is None) ^ (replacements_topic is None)
     supports_replacements = replacements_producer is not None
 
@@ -839,12 +839,13 @@ def flush_batch(
             inner_strategy = RunTaskWithMultiprocessing(
                 transform_function,
                 collect,
-                self.__processes,
                 max_batch_size=self.__max_batch_size,
                 max_batch_time=self.__max_batch_time,
+                pool=get_reusable_multiprocessing_pool(
+                    self.__processes, self.__initialize_parallel_transform
+                ),
                 input_block_size=self.__input_block_size,
                 output_block_size=self.__output_block_size,
-                initializer=self.__initialize_parallel_transform,
             )
 
         return RunTask(

snuba/consumers/strategy_factory.py

@@ -18,6 +18,7 @@
 
 from snuba.consumers.consumer import BytesInsertBatch, ProcessedMessageBatchWriter
 from snuba.consumers.dlq import ExitAfterNMessages
+from snuba.consumers.utils import get_reusable_multiprocessing_pool
 from snuba.processor import ReplacementBatch
 
 ProcessedMessage = Union[None, BytesInsertBatch, ReplacementBatch]
@@ -151,12 +152,13 @@ def flush_batch(
             strategy = RunTaskWithMultiprocessing(
                 transform_function,
                 collect,
-                self.__processes,
                 max_batch_size=self.__max_batch_size,
                 max_batch_time=self.__max_batch_time,
+                pool=get_reusable_multiprocessing_pool(
+                    self.__processes, self.__initialize_parallel_transform
+                ),
                 input_block_size=self.__input_block_size,
                 output_block_size=self.__output_block_size,
-                initializer=self.__initialize_parallel_transform,
             )
 
         if self.__prefilter is not None:

snuba/consumers/utils.py

@@ -1,6 +1,12 @@
+import atexit
 import logging
 import time
+from functools import lru_cache
+from typing import Callable, Optional
 
+from arroyo.processing.strategies.run_task_with_multiprocessing import (
+    MultiprocessingPool,
+)
 from confluent_kafka import KafkaException
 from confluent_kafka.admin import AdminClient
 
@@ -45,3 +51,17 @@
     total_partition_count = len(topic_metadata.partitions.keys())
     logger.info(f"Total of {total_partition_count} partition(s) for {topic.value}.")
     return total_partition_count
+
+
+@lru_cache(maxsize=None)
+def get_reusable_multiprocessing_pool(
+    num_processes: int, initializer: Optional[Callable[[], None]]
+) -> MultiprocessingPool:
+    pool = MultiprocessingPool(num_processes, initializer)
+
+    def shutdown() -> None:
+        logger.info("Shutting down multiprocessing pool")
+        pool.close()
+
+    atexit.register(shutdown)
+    return pool