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

@@ -39,6 +39,9 @@
     RunTaskInThreads,
     RunTaskWithMultiprocessing,
 )
+from arroyo.processing.strategies.run_task_with_multiprocessing import (
+    MultiprocessingPool,
+)
 from arroyo.types import (
     BaseValue,
     BrokerValue,
@@ -389,7 +392,6 @@
     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
 
@@ -800,6 +802,11 @@
             replacements,
             slice_id,
         )
+        self.__pool = (
+            MultiprocessingPool(self.__processes, initialize_parallel_transform)
+            if self.__processes is not None
+            else None
+        )
 
     def create_with_partitions(
         self,
@@ -833,18 +840,17 @@
             Union[FilteredPayload, MultistorageKafkaPayload]
         ]
 
-        if self.__processes is None:
+        if self.__pool is None:
             inner_strategy = RunTask(transform_function, collect)
         else:
             inner_strategy = RunTaskWithMultiprocessing(
                 transform_function,
                 collect,
-                self.__processes,
                 max_batch_size=self.__max_batch_size,
                 max_batch_time=self.__max_batch_time,
+                pool=self.__pool,
                 input_block_size=self.__input_block_size,
                 output_block_size=self.__output_block_size,
-                initializer=self.__initialize_parallel_transform,
             )
 
         return RunTask(
@@ -854,3 +860,7 @@
                 inner_strategy,
             ),
         )
+
+    def shutdown(self) -> None:
+        if self.__pool:
+            self.__pool.close()

snuba/consumers/strategy_factory.py

@@ -14,6 +14,9 @@
 )
 from arroyo.processing.strategies.commit import CommitOffsets
 from arroyo.processing.strategies.healthcheck import Healthcheck
+from arroyo.processing.strategies.run_task_with_multiprocessing import (
+    MultiprocessingPool,
+)
 from arroyo.types import BaseValue, Commit, FilteredPayload, Message, Partition
 
 from snuba.consumers.consumer import BytesInsertBatch, ProcessedMessageBatchWriter
@@ -98,6 +101,11 @@
         self.__output_block_size = output_block_size
         self.__initialize_parallel_transform = initialize_parallel_transform
         self.__health_check_file = health_check_file
+        self.__pool = (
+            MultiprocessingPool(self.__processes, self.__initialize_parallel_transform)
+            if self.__processes
+            else None
+        )
 
     def __should_accept(self, message: Message[KafkaPayload]) -> bool:
         assert self.__prefilter is not None
@@ -145,18 +153,17 @@
         transform_function = self.__process_message
 
         strategy: ProcessingStrategy[Union[FilteredPayload, KafkaPayload]]
-        if self.__processes is None:
+        if self.__pool is None:
             strategy = RunTask(transform_function, collect)
         else:
             strategy = RunTaskWithMultiprocessing(
                 transform_function,
                 collect,
-                self.__processes,
                 max_batch_size=self.__max_batch_size,
                 max_batch_time=self.__max_batch_time,
+                pool=self.__pool,
                 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:
@@ -173,3 +180,7 @@
             strategy = Healthcheck(self.__health_check_file, strategy)
 
         return strategy
+
+    def shutdown(self) -> None:
+        if self.__pool:
+            self.__pool.close()