diff --git a/ktables/kafka_table.py b/ktables/kafka_table.py index a3f69f6..bbde99b 100644 --- a/ktables/kafka_table.py +++ b/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,6 +862,11 @@ 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: @@ -843,6 +874,14 @@ async def stop(self) -> None: 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: diff --git a/tests/test_kafka_table.py b/tests/test_kafka_table.py index b36275d..e1645c7 100644 --- a/tests/test_kafka_table.py +++ b/tests/test_kafka_table.py @@ -257,6 +257,456 @@ async def test_barrier_after_stop_raises_stopped_not_unstarted(self) -> None: await table.barrier() +def _fake_consumer( + created: list, + *, + fail_start: bool = False, + fail_seek: bool = False, + partitions: tuple[int, ...] = (0,), + end_offset: int = 0, + catch_up: bool = True, + getmany_error: BaseException | None = None, + stop_error: BaseException | None = None, +) -> type: + """A configurable stand-in AIOKafkaConsumer for broker-free lifecycle tests. + + Every constructed instance is appended to ``created`` and records whether + ``stop()`` was awaited — so a test can assert on the exact consumer that + ``KafkaTable.start()`` built. + + - ``fail_start``: ``start()`` raises as if the broker were unreachable. + - ``fail_seek``: ``seek_to_beginning()`` raises (a post-connect setup failure). + - ``partitions``: partition numbers reported by ``assignment()``; ``()`` = none. + - ``end_offset``: end offset for every assigned partition. With the default 0, + the reader is caught up on its first poll (position 0 >= end offset 0). + - ``catch_up``: when False, ``position()`` stays at 0 so the reader never + latches — drives the degraded (catch-up-timeout) and cancellation paths. + Use with ``end_offset >= 1`` (0 < end_offset keeps the gate open). + - ``getmany_error``: if set, the reader loop's ``getmany()`` raises it — a + reader death. Combine with a non-zero ``end_offset`` to die *during* catch-up. + - ``stop_error``: if set, ``stop()`` raises it (after recording the attempt) — + to prove a teardown failure never masks the caller's original error. + """ + + class _FakeConsumer: + def __init__(self, *args: object, **kwargs: object) -> None: + self.stopped = False + topic = args[0] if args else kwargs.get("topic", "t") + self._tps = {TopicPartition(str(topic), p) for p in partitions} + self._fetch_max_wait_ms = kwargs.get("fetch_max_wait_ms") + created.append(self) + + async def start(self) -> None: + if fail_start: + raise KafkaConnectionError("simulated: broker unreachable") + + def assignment(self) -> set: + return set(self._tps) + + async def seek_to_beginning(self, *tps: object) -> None: + if fail_seek: + raise KafkaConnectionError("simulated: seek failed") + + async def end_offsets(self, tps: list) -> dict: + return {tp: end_offset for tp in tps} + + async def getmany(self, timeout_ms: int | None = None) -> dict: + if getmany_error is not None: + raise getmany_error + await asyncio.sleep(0.01) # yield; avoids a hot loop once caught up + return {} + + async def position(self, tp: object) -> int: + # Caught up once position reaches the end offset; catch_up=False pins + # it below so the reader loops indefinitely (degraded/cancellation). + return end_offset if catch_up else 0 + + async def stop(self) -> None: + self.stopped = True + if stop_error is not None: + raise stop_error + + return _FakeConsumer + + +def _fake_producer( + created: list, *, fail_start: bool = False, stop_error: BaseException | None = None +) -> type: + """A stand-in AIOKafkaProducer that records its instances and whether + ``stop()`` was awaited. ``fail_start`` makes ``start()`` raise as if the + broker were unreachable; ``stop_error`` makes ``stop()`` raise (after + recording the attempt) to prove rollback never masks the connect error.""" + + class _FakeProducer: + def __init__(self, **kwargs: object) -> None: + self.stopped = False + created.append(self) + + async def start(self) -> None: + if fail_start: + raise KafkaConnectionError("simulated: broker unreachable") + + async def stop(self) -> None: + self.stopped = True + if stop_error is not None: + raise stop_error + + return _FakeProducer + + +class TestStartFailureCleanup: + """start()/writer.start() must close the consumer/producer they construct + when startup fails. Otherwise the resource is orphaned: startup raises before + ``self._consumer``/``self._producer`` is assigned, so a follow-up ``stop()`` + has no reference to it and closes nothing — aiokafka then logs + ``Unclosed AIOKafka*`` (and leaks the socket) when it is garbage-collected. + + Regression tests for issue #30. No broker needed: the fakes fail + deterministically, so these pin the ktables-side contract ("the resource + start() created gets stopped on failure") without depending on GC timing or + aiokafka's ``__del__`` warning. + """ + + async def test_start_stops_consumer_when_initial_connect_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Issue #30 exact case: the initial broker connect (consumer.start()) + # raises, before self._consumer is assigned. + created: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(created, fail_start=True)) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.connect.fail", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=0.3, + ) + # The idiomatic caller: try/finally around start()/stop(). The connect + # failure must propagate out of start(), and the consumer start() built + # must be stopped along the way. + with pytest.raises(KafkaConnectionError): + try: + await table.start() + finally: + await table.stop() + + assert len(created) == 1, "start() should construct exactly one consumer" + assert created[0].stopped, ( + "start() orphaned the consumer it created: the initial connect failed but " + "the consumer was never stopped (self._consumer is still None, so stop() " + "closed nothing). See issue #30." + ) + # A failed initial connect leaves the table pristine: not started, so a + # retry is not blocked by an "already started" guard. + assert table.started is False + assert table.status == "unstarted" + assert table._consumer is None + + async def test_start_stops_consumer_when_no_partitions_assigned( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # A post-connect setup failure: connect succeeds but no partitions are + # assigned, so start() raises RuntimeError. The consumer must still be + # cleaned up (the failure guard must span the whole acquisition, not + # just the connect). + created: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(created, partitions=())) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.no.partitions", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=0.3, + ) + with pytest.raises(RuntimeError, match="no partitions"): + try: + await table.start() + finally: + await table.stop() + + assert created[0].stopped, "a post-connect setup failure must still stop the consumer" + assert table._consumer is None + assert table.started is False + + async def test_start_cleans_up_when_reader_dies_during_catchup( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The reader dies before catch-up completes: start() must raise, tear + # down BOTH the consumer and the reader task, and leave the table in the + # documented "failed but counts-as-started" state. + created: list = [] + monkeypatch.setattr( + "ktables.kafka_table.AIOKafkaConsumer", + _fake_consumer(created, end_offset=5, getmany_error=RuntimeError("induced reader death")), + ) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.reader.dies.during.catchup", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=1.0, + ) + with pytest.raises(RuntimeError, match="died during catch-up"): + try: + await table.start() + finally: + await table.stop() + + assert created[0].stopped, "reader death during catch-up must stop the consumer" + assert table._consumer is None + assert table._task is None + assert table.started is True, "after a reader-death failure the table counts as started" + assert table.status == "failed" + + async def test_successful_start_commits_consumer_without_stopping_it( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The commit path: on success the consumer must NOT be stopped (the + # cleanup guard is disarmed), the table reaches caught_up, and only an + # explicit stop() tears it down. + created: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(created)) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.success", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=1.0, + ) + await table.start() + try: + assert table.status == "caught_up" + assert table.started is True + assert table._consumer is created[0] + assert table._task is not None + assert created[0].stopped is False, "a successful start() must not stop the consumer it committed" + finally: + await table.stop() + + assert created[0].stopped is True, "stop() must tear the committed consumer down" + assert table._consumer is None + assert table._task is None + + async def test_writer_start_stops_producer_when_initial_connect_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The writer has the same shape as the consumer path: producer.start() + # raises before self._producer is assigned, orphaning the producer. + created: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaProducer", _fake_producer(created, fail_start=True)) + writer = KafkaTableWriter.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.writer.connect.fail", + model=ServiceRecord, + ensure_topic=False, + ) + with pytest.raises(KafkaConnectionError): + try: + await writer.start() + finally: + await writer.stop() + + assert len(created) == 1, "writer.start() should construct exactly one producer" + assert created[0].stopped, ( + "writer.start() orphaned the producer it created: the initial connect failed but " + "the producer was never stopped (self._producer is still None). See issue #30." + ) + + async def test_start_as_context_manager_stops_consumer_when_connect_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The real issue #30 trigger, and the one that ISOLATES start()'s own + # rollback: `async with table:` calls start() from __aenter__, and when + # start() raises, Python does NOT call __aexit__ — so no stop() ever + # runs. The consumer can therefore only be closed by the AsyncExitStack + # inside start() itself (unlike the try/finally tests, where a trailing + # stop() could mask a broken internal rollback). + created: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(created, fail_start=True)) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.async.with.connect.fail", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=0.3, + ) + with pytest.raises(KafkaConnectionError): + async with table: # __aenter__ raises; body and __aexit__ never run + pass + + assert len(created) == 1 + assert created[0].stopped, ( + "start() failed to close the consumer on its own: with `async with`, no stop() " + "runs after __aenter__ fails, so only start()'s internal rollback can close it. " + "See issue #30." + ) + assert table._consumer is None + assert table.started is False + + async def test_start_stops_consumer_when_seek_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Another post-connect setup failure (seek_to_beginning raises): the + # AsyncExitStack must still unwind the consumer. Exercised via `async + # with` so only start()'s internal rollback can satisfy the assertion. + created: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(created, fail_seek=True)) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.seek.fail", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=0.3, + ) + with pytest.raises(KafkaConnectionError, match="seek"): + async with table: + pass + + assert created[0].stopped, "a seek failure must still stop the consumer" + assert table._consumer is None + assert table.started is False + + async def test_retry_after_failed_connect_succeeds( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # A failed initial connect must leave the table retriable — not wedged + # behind the "already started" guard. First attempt fails; a second, + # against a now-reachable broker, starts cleanly. + failing: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(failing, fail_start=True)) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.retry", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=1.0, + ) + with pytest.raises(KafkaConnectionError): + await table.start() + + succeeding: list = [] + monkeypatch.setattr("ktables.kafka_table.AIOKafkaConsumer", _fake_consumer(succeeding)) + await table.start() + try: + assert table.status == "caught_up" + assert table._consumer is succeeding[0] + finally: + await table.stop() + + async def test_start_surfaces_connect_error_even_if_consumer_stop_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Rollback must not mask the caller's error: if the consumer's own stop() + # ALSO fails during unwind, start() must still raise the ORIGINAL connect + # error, not the teardown error. _detach_consumer logs-and-swallows the + # stop() failure to guarantee this. + created: list = [] + monkeypatch.setattr( + "ktables.kafka_table.AIOKafkaConsumer", + _fake_consumer(created, fail_start=True, stop_error=RuntimeError("stop boom")), + ) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.stop.also.fails", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=0.3, + ) + with pytest.raises(KafkaConnectionError): + async with table: + pass + assert created[0].stopped, "stop() should have been attempted during rollback" + + async def test_start_cancelled_during_catchup_tears_down_consumer_and_task( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Cancelling an in-flight start() parked in catch-up must unwind the + # AsyncExitStack — tearing down BOTH the reader task and the consumer. + # This is the only path that pins the _detach_task stack registration: a + # LIVE task, not an already-dead one (as in the reader-death test). + created: list = [] + monkeypatch.setattr( + "ktables.kafka_table.AIOKafkaConsumer", + _fake_consumer(created, end_offset=1, catch_up=False), # never latches + ) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.cancelled.during.catchup", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=5.0, # long — we cancel well before it elapses + ) + starting = asyncio.create_task(table.start()) + # Wait until start() created the reader task and parked in catch-up. + assert await eventually(lambda: table._task is not None) + await asyncio.sleep(0.05) + + starting.cancel() + with pytest.raises(asyncio.CancelledError): + await starting + + assert created[0].stopped, "cancellation must stop the consumer" + assert table._consumer is None, "cancellation must detach the consumer" + assert table._task is None, "cancellation must cancel and detach the reader task" + + async def test_start_degrades_on_catchup_timeout_and_keeps_serving( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Catch-up timeout is a COMMIT, not a failure: start() must still reach + # pop_all(), so the consumer and reader task stay alive and keep serving + # (status "degraded"). A regression that tore them down on timeout would + # STILL read status "degraded" (derived from _timed_out), so assert on + # the live resources directly. + created: list = [] + monkeypatch.setattr( + "ktables.kafka_table.AIOKafkaConsumer", + _fake_consumer(created, end_offset=1, catch_up=False), # never latches + ) + table = KafkaTable.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.start.degraded", + model=ServiceRecord, + ensure_topic=False, + catchup_timeout=0.2, # times out fast -> degraded + ) + await table.start() + try: + assert table.status == "degraded" + assert table.started is True + assert table._consumer is created[0] + assert table._task is not None + assert created[0].stopped is False, ( + "a degraded (timed-out) start must COMMIT via pop_all() and keep the consumer " + "alive — the table serves partial data, it does not tear down" + ) + finally: + await table.stop() + assert created[0].stopped is True + assert table._consumer is None + assert table._task is None + + async def test_writer_start_surfaces_connect_error_even_if_producer_stop_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Parity with the consumer's _detach_consumer: if producer.stop() also + # fails during rollback, the ORIGINAL KafkaConnectionError must still + # propagate, not the teardown error. + created: list = [] + monkeypatch.setattr( + "ktables.kafka_table.AIOKafkaProducer", + _fake_producer(created, fail_start=True, stop_error=RuntimeError("stop boom")), + ) + writer = KafkaTableWriter.json( + bootstrap_servers=BOOTSTRAP, + topic="unit.writer.stop.also.fails", + model=ServiceRecord, + ensure_topic=False, + ) + with pytest.raises(KafkaConnectionError): + await writer.start() + assert created[0].stopped, "stop() should have been attempted during rollback" + + # --------------------------------------------------------------------------- # Integration tests (broker required; Redpanda auto-started via testcontainers) # --------------------------------------------------------------------------- diff --git a/uv.lock b/uv.lock index 61b6fc1..21cc7cc 100644 --- a/uv.lock +++ b/uv.lock @@ -409,7 +409,7 @@ wheels = [ [[package]] name = "ktables" -version = "0.3.0" +version = "1.1.0" source = { editable = "." } dependencies = [ { name = "aiokafka" },