fix: close consumer and producer when initial broker connect fails

ktables/kafka_table.py

@@ -42,6 +42,7 @@
 import logging
 import time
 from collections.abc import Callable, Iterator, Mapping
+from contextlib import AsyncExitStack
 from dataclasses import dataclass
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar
@@ -753,21 +754,35 @@ async def start(self) -> None:
                 topic_configs=self._topic_configs,
                 on_policy_mismatch=self._on_policy_mismatch,
             )
-        # Constructor topic + group_id=None: all partitions are assigned
-        # synchronously inside start() — no group, no commits, no rebalance.
-        # (Manual assign()+partitions_for_topic() hits a stale partition
-        # cache on fresh topics.)
-        consumer = AIOKafkaConsumer(
-            self._topic,
-            bootstrap_servers=self._bootstrap_servers,
-            group_id=None,
-            enable_auto_commit=False,
-            auto_offset_reset="earliest",
-            fetch_max_wait_ms=self._fetch_max_wait_ms,
-        )
-        await consumer.start()
-        self._consumer = consumer
-        try:
+        # A single AsyncExitStack owns everything start() acquires — the consumer
+        # and the reader task. Each cleanup is registered the instant its resource
+        # exists (BEFORE the fallible connect), so any failure or cancellation
+        # unwinds them in LIFO order (task, then consumer). On full success
+        # pop_all() disarms the stack and ownership transfers to the instance.
+        # This makes an orphaned, never-stopped consumer structurally impossible
+        # (issue #30). Note self._consumer is assigned BEFORE consumer.start()
+        # (below), on purpose: the registered _detach_consumer must run against a
+        # live handle when the connect fails. (The original bug assigned it only
+        # after a successful connect, so a caller's stop() — or, via `async with`,
+        # no stop() at all — had nothing to close.) The registered callbacks are
+        # the exact steps stop() runs, so start()'s rollback and stop() never
+        # diverge.
+        async with AsyncExitStack() as stack:
+            # Constructor topic + group_id=None: all partitions are assigned
+            # synchronously inside consumer.start() — no group, no commits, no
+            # rebalance. (Manual assign()+partitions_for_topic() hits a stale
+            # partition cache on fresh topics.)
+            consumer = AIOKafkaConsumer(
+                self._topic,
+                bootstrap_servers=self._bootstrap_servers,
+                group_id=None,
+                enable_auto_commit=False,
+                auto_offset_reset="earliest",
+                fetch_max_wait_ms=self._fetch_max_wait_ms,
+            )
+            self._consumer = consumer
+            stack.push_async_callback(self._detach_consumer)
+            await consumer.start()
             tps = sorted(consumer.assignment(), key=lambda tp: tp.partition)
             if not tps:
                 raise RuntimeError(
@@ -779,27 +794,30 @@ async def start(self) -> None:
             await consumer.seek_to_beginning(*tps)
             # Gate target: HWM at start; later records are live updates.
             end_offsets: dict[TopicPartition, int] = await consumer.end_offsets(tps)
-        except BaseException:
-            await consumer.stop()
-            self._consumer = None
-            raise
-        self._started = True
-        task = asyncio.create_task(self._run(consumer, tps, end_offsets, time.perf_counter()), name=f"kafka-table:{self._topic}")
-        task.add_done_callback(self._on_reader_done)
-        self._task = task
-        if not await self.wait_until_caught_up(timeout=self._catchup_timeout):
-            if self._failure is not None:
-                # Reader died during boot: fail start() loudly, cleaned up.
-                failure = self._failure
-                await self.stop()
-                raise RuntimeError(f"KafkaTable reader for topic={self._topic!r} died during catch-up") from failure
-            self._timed_out = True
-            logger.error(
-                "table for topic=%s NOT caught up after %.1fs (applied=%d so far); continuing DEGRADED — data may be incomplete",
-                self._topic,
-                self._catchup_timeout,
-                self._live.records_applied,
-            )
+
+            self._started = True
+            task = asyncio.create_task(self._run(consumer, tps, end_offsets, time.perf_counter()), name=f"kafka-table:{self._topic}")
+            task.add_done_callback(self._on_reader_done)
+            self._task = task
+            stack.push_async_callback(self._detach_task)
+
+            if not await self.wait_until_caught_up(timeout=self._catchup_timeout):
+                if self._failure is not None:
+                    # Reader died during catch-up: let the stack tear down the
+                    # task and consumer, then fail start() loudly. self._started
+                    # stays True — the table counts as started (status "failed").
+                    raise RuntimeError(f"KafkaTable reader for topic={self._topic!r} died during catch-up") from self._failure
+                self._timed_out = True
+                logger.error(
+                    "table for topic=%s NOT caught up after %.1fs (applied=%d so far); continuing DEGRADED — data may be incomplete",
+                    self._topic,
+                    self._catchup_timeout,
+                    self._live.records_applied,
+                )
+            # Caught up or degraded (both keep the table serving): commit —
+            # ownership of the consumer and reader task passes to the instance
+            # and the stack unwinds nothing.
+            stack.pop_all()
 
     def _on_reader_done(self, task: asyncio.Task[None]) -> None:
         if task.cancelled():
@@ -816,7 +834,15 @@ def _on_reader_done(self, task: asyncio.Task[None]) -> None:
                 exc_info=exc,
             )
 
-    async def stop(self) -> None:
+    async def _detach_task(self) -> None:
+        """Cancel and join the reader task, then wake any pending barriers.
+
+        Shared by ``stop()`` and ``start()``'s failure rollback (registered on
+        the AsyncExitStack), so the two teardown paths never diverge. Never
+        raises: the reader's own error was already captured by
+        ``_on_reader_done``, and re-raising here would mask the caller's error
+        and skip consumer cleanup.
+        """
         task, self._task = self._task, None
         if task is not None:
             task.cancel()
@@ -836,13 +862,26 @@ async def stop(self) -> None:
             if not fut.done():
                 fut.cancel()
         self._barriers.clear()
+
+    async def _detach_consumer(self) -> None:
+        """Stop the consumer and drop the reference. Shared by ``stop()`` and
+        ``start()``'s failure rollback; never raises (a teardown-time
+        ``stop()`` error is logged, not propagated)."""
         consumer, self._consumer = self._consumer, None
         if consumer is not None:
             try:
                 await consumer.stop()
             except Exception:
                 logger.exception("consumer.stop() failed for topic=%s during teardown", self._topic)
 
+    async def stop(self) -> None:
+        # LIFO teardown: task first (so it stops touching the consumer), then
+        # the consumer. start()'s rollback registers these same two callbacks on
+        # its AsyncExitStack in reverse push order (consumer, then task) so the
+        # LIFO unwind runs them task-first — matching this order exactly.
+        await self._detach_task()
+        await self._detach_consumer()
+
     async def __aenter__(self) -> Self:
         await self.start()
         return self
@@ -1044,7 +1083,22 @@ async def start(self) -> None:
             enable_idempotence=self._enable_idempotence,
             **optional_kwargs,
         )
-        await producer.start()
+        # Close the producer if the initial connect fails, so a failed start()
+        # never orphans it — otherwise aiokafka logs "Unclosed AIOKafkaProducer"
+        # and leaks the socket (issue #30). stop() cannot help: self._producer is
+        # only assigned on success. Single resource, single step — a plain guard
+        # is the right weight here (cf. KafkaTable.start()'s AsyncExitStack).
+        try:
+            await producer.start()
+        except BaseException:
+            # Guard stop() the same way _detach_consumer does: a teardown failure
+            # must not mask the caller's connect error — the bare `raise` below
+            # must re-raise the ORIGINAL error, not a stop() error.
+            try:
+                await producer.stop()
+            except Exception:
+                logger.exception("producer.stop() failed for topic=%s during failed-start rollback", self._topic)
+            raise
         self._producer = producer
 
     async def stop(self) -> None: