buz 2.15.10rc1__tar.gz → 2.15.10rc3__tar.gz

{buz-2.15.10rc1 → buz-2.15.10rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: buz
-Version: 2.15.10rc1
+Version: 2.15.10rc3
 Summary: Buz is a set of light, simple and extensible implementations of event, command and query buses.
 License: MIT
 Author: Luis Pintado Lozano

{buz-2.15.10rc1 → buz-2.15.10rc3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "buz"
-version = "2.15.10-rc1"
+version = "2.15.10-rc3"
 description = "Buz is a set of light, simple and extensible implementations of event, command and query buses."
 readme = "README.md"
 authors = ["Luis Pintado Lozano <luis.pintado.lozano@gmail.com>", "Gerardo Parra <gprauxiliar@gmail.com>"]

{buz-2.15.10rc1 → buz-2.15.10rc3}/src/buz/event/infrastructure/buz_kafka/base_buz_aiokafka_async_consumer.py

@@ -226,6 +226,7 @@ class BaseBuzAIOKafkaAsyncConsumer(AsyncConsumer):
             logger=self._logger,
             kafka_admin_client=self.__kafka_admin_client,
             auto_create_topic_configuration=self.__auto_create_topic_configuration,
+            on_partition_assigned=self.__on_partition_assigned,
             on_partition_revoked=self.__on_partition_revoked,
             session_timeout_ms=self.__session_timeout_ms,
             max_poll_interval_ms=self.__max_poll_interval_ms,
@@ -328,36 +329,41 @@ class BaseBuzAIOKafkaAsyncConsumer(AsyncConsumer):
 
             while consumer != last_consumer:
                 consumer, queue = next(consumer_queues_cyclic_iterator)
-                kafka_poll_record = queue.pop()
+                async with self.__get_consumer_mutex(
+                    consumer_fqn=consumer.get_consumer_group(),
+                ):
+                    kafka_poll_record = queue.pop()
 
-                if kafka_poll_record is not None:
-                    consuming_task = ConsumingTask(consumer, kafka_poll_record)
-                    async with self.__get_partition_mutex(
-                        consumer_fqn=consuming_task.consumer.get_consumer_group(),
-                        partition=kafka_poll_record.partition,
-                    ):
-                        yield consuming_task
+                    if kafka_poll_record is not None:
+                        yield ConsumingTask(consumer, kafka_poll_record)
 
                     last_consumer = consumer
                     break
 
-    def __get_partition_mutex(self, consumer_fqn: str, partition: int) -> Lock:
-        mutex_key = f"consumer_{consumer_fqn}_partition_{partition}"
+    def __get_consumer_mutex(self, consumer_fqn: str) -> Lock:
+        mutex_key = f"consumer_{consumer_fqn}"
         return self.__consumer_and_partition_mutex[mutex_key]
 
     async def __all_queues_are_empty(self) -> bool:
         return all([queue.is_totally_empty() for queue in self.__queue_per_consumer_mapper.values()])
 
+    async def __on_partition_assigned(self, consumer: AIOKafkaConsumer, topics_partitions: set[TopicPartition]) -> None:
+        self._logger.info(
+            f"rebalancing in progress, assigning partitions {topics_partitions} to consumer {consumer.get_consumer_group()}"
+        )
+
     async def __on_partition_revoked(self, consumer: AIOKafkaConsumer, topics_partitions: set[TopicPartition]) -> None:
-        self._logger.info(f"rebalancing in progress, revoking partitions {topics_partitions}")
+        self._logger.info(
+            f"rebalancing in progress, revoking partitions {topics_partitions} from consumer {consumer.get_consumer_group()}"
+        )
 
         self.increase_number_of_rebalancing_in_progress()
-        for topic_partition in topics_partitions:
-            async with self.__get_partition_mutex(
-                consumer_fqn=consumer.get_consumer_group(),
-                partition=topic_partition.partition,
-            ):
+        async with self.__get_consumer_mutex(
+            consumer_fqn=consumer.get_consumer_group(),
+        ):
+            for topic_partition in topics_partitions:
                 self.__queue_per_consumer_mapper[consumer].clear(topic_partition)
+
         self.decrease_number_of_rebalancing_in_progress()
 
     def increase_number_of_rebalancing_in_progress(self) -> None:
@@ -374,8 +380,7 @@ class BaseBuzAIOKafkaAsyncConsumer(AsyncConsumer):
         self._logger.info("Worker stop requested. Waiting for finalize the current task")
 
     async def __manage_kafka_consumers_stopping(self) -> None:
-        for kafka_consumer in self.__queue_per_consumer_mapper.keys():
-            await kafka_consumer.stop()
+        await gather(*[kafka_consumer.stop() for kafka_consumer in self.__queue_per_consumer_mapper.keys()])
 
     async def __health_check(self) -> web.Response:
         health_information = {

{buz-2.15.10rc1 → buz-2.15.10rc3}/src/buz/kafka/infrastructure/aiokafka/aiokafka_consumer.py

@@ -4,7 +4,7 @@ import asyncio
 from logging import Logger
 from typing import Awaitable, Callable, Optional, Sequence, cast
 
-from aiokafka import AIOKafkaConsumer as AIOKafkaNativeConsumer, TopicPartition, OffsetAndMetadata
+from aiokafka import AIOKafkaConsumer as AIOKafkaNativeConsumer, ConsumerRecord, TopicPartition, OffsetAndMetadata
 from aiokafka.helpers import create_ssl_context
 
 from buz.event.infrastructure.buz_kafka.exceptions.kafka_event_bus_config_not_valid_exception import (
@@ -26,6 +26,7 @@ from buz.kafka.infrastructure.aiokafka.translators.consumer_initial_offset_posit
 
 class AIOKafkaConsumer:
     __DEFAULT_POLL_TIMEOUT_MS = 0
+    __DEFAULT_ASYNC_FUNCTION_TIMEOUT_MS = 3000
 
     def __init__(
         self,
@@ -41,6 +42,7 @@ class AIOKafkaConsumer:
         max_poll_interval_ms: int,
         heartbeat_interval_ms: int,
         on_partition_revoked: Callable[[AIOKafkaConsumer, set[TopicPartition]], Awaitable[None]],
+        on_partition_assigned: Callable[[AIOKafkaConsumer, set[TopicPartition]], Awaitable[None]],
         auto_create_topic_configuration: Optional[AutoCreateTopicConfiguration] = None,
         wait_for_connection_to_cluster_ms: Optional[int] = None,
         group_instance_id: Optional[str] = None,
@@ -55,6 +57,7 @@ class AIOKafkaConsumer:
         self.__session_timeout_ms = session_timeout_ms
         self.__auto_create_topic_configuration = auto_create_topic_configuration
         self.__on_partitions_revoked_callback = on_partition_revoked
+        self.__on_partitions_assigned_callback = on_partition_assigned
         self.__max_poll_interval_ms = max_poll_interval_ms
         self.__heartbeat_interval_ms = heartbeat_interval_ms
         self.__wait_for_connection_to_cluster_ms = wait_for_connection_to_cluster_ms
@@ -166,7 +169,7 @@ class AIOKafkaConsumer:
         self,
         topics_partitions: set[TopicPartition],
     ) -> None:
-        return
+        await self.__on_partitions_assigned_callback(self, topics_partitions)
 
     async def __on_partitions_revoked(
         self,
@@ -177,14 +180,15 @@ class AIOKafkaConsumer:
     async def poll(
         self,
         *,
-        timeout_ms: int = __DEFAULT_POLL_TIMEOUT_MS,
         number_of_messages_to_poll: Optional[int] = None,
     ) -> list[KafkaPollRecord]:
-        poll_results = await self.__consumer.getmany(
-            timeout_ms=timeout_ms,
+        poll_results = await self.__get_many(
             max_records=number_of_messages_to_poll,
         )
 
+        if poll_results is None:
+            return []
+
         results = [
             cast(KafkaPollRecord, consumer_record)
             for consumer_records in poll_results.values()
@@ -193,6 +197,19 @@ class AIOKafkaConsumer:
 
         return results
 
+    async def __get_many(
+        self,
+        max_records: Optional[int] = None,
+    ) -> Optional[dict[TopicPartition, list[ConsumerRecord]]]:
+        try:
+            return await asyncio.wait_for(
+                self.__consumer.getmany(timeout_ms=self.__DEFAULT_POLL_TIMEOUT_MS, max_records=max_records),
+                timeout=self.__DEFAULT_ASYNC_FUNCTION_TIMEOUT_MS / 1000,
+            )
+        except asyncio.TimeoutError:
+            self.__logger.debug("Timeout while polling")
+            return None
+
     async def commit_poll_record(self, poll_record: KafkaPollRecord) -> None:
         topic_partition = TopicPartition(topic=poll_record.topic, partition=poll_record.partition)
         offset = {topic_partition: OffsetAndMetadata(poll_record.offset + 1, "")}
@@ -201,4 +218,13 @@ class AIOKafkaConsumer:
 
     async def stop(self) -> None:
         self.__logger.info(f"Closing connection of consumer with group_id={self.__consumer_group}")
-        await self.__consumer.stop()
+
+        self.__consumer.unsubscribe()
+
+        try:
+            await asyncio.wait_for(
+                self.__consumer.stop(),
+                timeout=self.__DEFAULT_ASYNC_FUNCTION_TIMEOUT_MS / 1000,
+            )
+        except asyncio.TimeoutError:
+            self.__logger.debug("Timeout while stopping consumer")