Handle unclosed producer and consumer warnings (#1130)

Signed-off-by: Callan Gray <[email protected]>

Author: calgray

CHANGES.rst

@@ -16,6 +16,8 @@ Bugfixes:
 
 * Make KafkaStorageError retriable after metadata refresh like in other
   implementations (pr #1115 by @omerhadari)
+* Fix producer and consumer requiring closing after `start` or `__aenter__` raise an exception.
+  (issue #1130, pr #1131 by @calgray)
 
 
 Misc:

aiokafka/consumer/consumer.py

@@ -4,6 +4,7 @@
 import sys
 import traceback
 import warnings
+from contextlib import AsyncExitStack
 
 from aiokafka import __version__
 from aiokafka.abc import ConsumerRebalanceListener
@@ -335,7 +336,6 @@ def __init__(
 
         if loop.get_debug():
             self._source_traceback = traceback.extract_stack(sys._getframe(1))
-        self._closed = False
 
         if topics:
             topics = self._validate_topics(topics)
@@ -368,83 +368,98 @@ async def start(self):
             self._loop is get_running_loop()
         ), "Please create objects with the same loop as running with"
         assert self._fetcher is None, "Did you call `start` twice?"
-        await self._client.bootstrap()
-        await self._wait_topics()
 
-        if self._client.api_version < (0, 9):
-            raise ValueError(f"Unsupported Kafka version: {self._client.api_version}")
+        self._exit_stack = AsyncExitStack()
+        try:
+            await self._client.bootstrap()
+            self._exit_stack.push_async_callback(self._client.close)
+            await self._wait_topics()
 
-        if (
-            self._isolation_level == "read_committed"
-            and self._client.api_version < (0, 11)  # fmt: skip
-        ):
-            raise UnsupportedVersionError(
-                "`read_committed` isolation_level available only for Brokers "
-                "0.11 and above"
-            )
+            if self._client.api_version < (0, 9):
+                raise ValueError(
+                    f"Unsupported Kafka version: {self._client.api_version}"
+                )
 
-        self._fetcher = Fetcher(
-            self._client,
-            self._subscription,
-            key_deserializer=self._key_deserializer,
-            value_deserializer=self._value_deserializer,
-            fetch_min_bytes=self._fetch_min_bytes,
-            fetch_max_bytes=self._fetch_max_bytes,
-            fetch_max_wait_ms=self._fetch_max_wait_ms,
-            max_partition_fetch_bytes=self._max_partition_fetch_bytes,
-            check_crcs=self._check_crcs,
-            fetcher_timeout=self._consumer_timeout,
-            retry_backoff_ms=self._retry_backoff_ms,
-            auto_offset_reset=self._auto_offset_reset,
-            isolation_level=self._isolation_level,
-        )
+            if (
+                self._isolation_level == "read_committed"
+                and self._client.api_version < (0, 11)  # fmt: skip
+            ):
+                raise UnsupportedVersionError(
+                    "`read_committed` isolation_level available only for Brokers "
+                    "0.11 and above"
+                )
 
-        if self._group_id is not None:
-            # using group coordinator for automatic partitions assignment
-            self._coordinator = GroupCoordinator(
+            self._fetcher = Fetcher(
                 self._client,
                 self._subscription,
-                group_id=self._group_id,
-                group_instance_id=self._group_instance_id,
-                heartbeat_interval_ms=self._heartbeat_interval_ms,
-                session_timeout_ms=self._session_timeout_ms,
+                key_deserializer=self._key_deserializer,
+                value_deserializer=self._value_deserializer,
+                fetch_min_bytes=self._fetch_min_bytes,
+                fetch_max_bytes=self._fetch_max_bytes,
+                fetch_max_wait_ms=self._fetch_max_wait_ms,
+                max_partition_fetch_bytes=self._max_partition_fetch_bytes,
+                check_crcs=self._check_crcs,
+                fetcher_timeout=self._consumer_timeout,
                 retry_backoff_ms=self._retry_backoff_ms,
-                enable_auto_commit=self._enable_auto_commit,
-                auto_commit_interval_ms=self._auto_commit_interval_ms,
-                assignors=self._partition_assignment_strategy,
-                exclude_internal_topics=self._exclude_internal_topics,
-                rebalance_timeout_ms=self._rebalance_timeout_ms,
-                max_poll_interval_ms=self._max_poll_interval_ms,
+                auto_offset_reset=self._auto_offset_reset,
+                isolation_level=self._isolation_level,
             )
-            if self._subscription.subscription is not None:
-                if self._subscription.partitions_auto_assigned():
+            self._exit_stack.push_async_callback(self._fetcher.close)
+
+            if self._group_id is not None:
+                # using group coordinator for automatic partitions assignment
+                self._coordinator = GroupCoordinator(
+                    self._client,
+                    self._subscription,
+                    group_id=self._group_id,
+                    group_instance_id=self._group_instance_id,
+                    heartbeat_interval_ms=self._heartbeat_interval_ms,
+                    session_timeout_ms=self._session_timeout_ms,
+                    retry_backoff_ms=self._retry_backoff_ms,
+                    enable_auto_commit=self._enable_auto_commit,
+                    auto_commit_interval_ms=self._auto_commit_interval_ms,
+                    assignors=self._partition_assignment_strategy,
+                    exclude_internal_topics=self._exclude_internal_topics,
+                    rebalance_timeout_ms=self._rebalance_timeout_ms,
+                    max_poll_interval_ms=self._max_poll_interval_ms,
+                )
+                self._exit_stack.push_async_callback(self._coordinator.close)
+
+                if self._subscription.subscription is not None:
+                    if self._subscription.partitions_auto_assigned():
+                        # Either we passed `topics` to constructor or `subscribe`
+                        # was called before `start`
+                        await self._subscription.wait_for_assignment()
+                    else:
+                        # `assign` was called before `start`. We did not start
+                        # this task on that call, as coordinator was yet to be
+                        # created
+                        self._coordinator.start_commit_offsets_refresh_task(
+                            self._subscription.subscription.assignment
+                        )
+            else:
+                # Using a simple assignment coordinator for reassignment on
+                # metadata changes
+                self._coordinator = NoGroupCoordinator(
+                    self._client,
+                    self._subscription,
+                    exclude_internal_topics=self._exclude_internal_topics,
+                )
+                self._exit_stack.push_async_callback(self._coordinator.close)
+
+                if (
+                    self._subscription.subscription is not None
+                    and self._subscription.partitions_auto_assigned()
+                ):
                     # Either we passed `topics` to constructor or `subscribe`
                     # was called before `start`
-                    await self._subscription.wait_for_assignment()
-                else:
-                    # `assign` was called before `start`. We did not start
-                    # this task on that call, as coordinator was yet to be
-                    # created
-                    self._coordinator.start_commit_offsets_refresh_task(
-                        self._subscription.subscription.assignment
-                    )
-        else:
-            # Using a simple assignment coordinator for reassignment on
-            # metadata changes
-            self._coordinator = NoGroupCoordinator(
-                self._client,
-                self._subscription,
-                exclude_internal_topics=self._exclude_internal_topics,
-            )
+                    await self._client.force_metadata_update()
+                    self._coordinator.assign_all_partitions(check_unknown=True)
 
-            if (
-                self._subscription.subscription is not None
-                and self._subscription.partitions_auto_assigned()
-            ):
-                # Either we passed `topics` to constructor or `subscribe`
-                # was called before `start`
-                await self._client.force_metadata_update()
-                self._coordinator.assign_all_partitions(check_unknown=True)
+            self._closed = False
+        except:
+            await self._exit_stack.aclose()
+            raise
 
     async def _wait_topics(self):
         if self._subscription.subscription is not None:
@@ -514,11 +529,7 @@ async def stop(self):
             return
         log.debug("Closing the KafkaConsumer.")
         self._closed = True
-        if self._coordinator:
-            await self._coordinator.close()
-        if self._fetcher:
-            await self._fetcher.close()
-        await self._client.close()
+        await self._exit_stack.aclose()
         log.debug("The KafkaConsumer has closed.")
 
     async def commit(self, offsets=None):

aiokafka/producer/producer.py

@@ -3,6 +3,7 @@
 import sys
 import traceback
 import warnings
+from contextlib import AsyncExitStack
 
 from aiokafka.client import AIOKafkaClient
 from aiokafka.codec import has_gzip, has_lz4, has_snappy, has_zstd
@@ -324,8 +325,6 @@ def __init__(
             request_timeout_ms=request_timeout_ms,
         )
 
-        self._closed = False
-
     # Warn if producer was not closed properly
     # We don't attempt to close the Consumer, as __del__ is synchronous
     def __del__(self, _warnings=warnings):
@@ -348,28 +347,39 @@ async def start(self):
         assert (
             self._loop is get_running_loop()
         ), "Please create objects with the same loop as running with"
-        log.debug("Starting the Kafka producer")  # trace
-        await self.client.bootstrap()
-
-        if self._compression_type == "lz4":
-            assert self.client.api_version >= (0, 8, 2), (
-                "LZ4 Requires >= Kafka 0.8.2 Brokers"
-            )  # fmt: skip
-        elif self._compression_type == "zstd":
-            assert self.client.api_version >= (2, 1, 0), (
-                "Zstd Requires >= Kafka 2.1.0 Brokers"
-            )  # fmt: skip
-
-        if self._txn_manager is not None and self.client.api_version < (0, 11):
-            raise UnsupportedVersionError(
-                "Idempotent producer available only for Broker version 0.11"
-                " and above"
-            )
 
-        await self._sender.start()
-        self._message_accumulator.set_api_version(self.client.api_version)
-        self._producer_magic = 0 if self.client.api_version < (0, 10) else 1
-        log.debug("Kafka producer started")
+        self._exit_stack = AsyncExitStack()
+        try:
+            log.debug("Starting the Kafka producer")  # trace
+
+            await self.client.bootstrap()
+            self._exit_stack.push_async_callback(self.client.close)
+
+            if self._compression_type == "lz4":
+                assert self.client.api_version >= (0, 8, 2), (
+                    "LZ4 Requires >= Kafka 0.8.2 Brokers"
+                )  # fmt: skip
+            elif self._compression_type == "zstd":
+                assert self.client.api_version >= (2, 1, 0), (
+                    "Zstd Requires >= Kafka 2.1.0 Brokers"
+                )  # fmt: skip
+
+            if self._txn_manager is not None and self.client.api_version < (0, 11):
+                raise UnsupportedVersionError(
+                    "Idempotent producer available only for Broker version 0.11"
+                    " and above"
+                )
+
+            await self._sender.start()
+            self._exit_stack.push_async_callback(self._sender.close)
+
+            self._message_accumulator.set_api_version(self.client.api_version)
+            self._producer_magic = 0 if self.client.api_version < (0, 10) else 1
+            self._closed = False
+            log.debug("Kafka producer started")
+        except:
+            await self._exit_stack.aclose()
+            raise
 
     async def flush(self):
         """Wait until all batches are Delivered and futures resolved"""
@@ -379,6 +389,7 @@ async def stop(self):
         """Flush all pending data and close all connections to kafka cluster"""
         if self._closed:
             return
+        log.debug("Closing the KafkaProducer.")
         self._closed = True
 
         # If the sender task is down there is no way for accumulator to flush
@@ -391,9 +402,7 @@ async def stop(self):
                 return_when=asyncio.FIRST_COMPLETED,
             )
 
-            await self._sender.close()
-
-        await self.client.close()
+        await self._exit_stack.aclose()
         log.debug("The Kafka producer has closed.")
 
     async def partitions_for(self, topic):