From 3b489ca718d0685b215d667c8cfcf917ec9c5d0b Mon Sep 17 00:00:00 2001 From: William Barnhart Date: Wed, 21 Oct 2026 22:45:51 -0400 Subject: [PATCH] Track transactional offsets in aiokafka monitor --- faust/transport/consumer.py | 12 ++++++++++++ faust/transport/drivers/aiokafka.py | 9 +++++++++ tests/unit/transport/drivers/test_aiokafka.py | 15 +++++++++++++++ tests/unit/transport/test_consumer.py | 6 ++++++ 4 files changed, 42 insertions(+) diff --git a/faust/transport/consumer.py b/faust/transport/consumer.py index 9e49a9a8f..391422877 100644 --- a/faust/transport/consumer.py +++ b/faust/transport/consumer.py @@ -1069,6 +1069,8 @@ async def _commit_offsets( committable_offsets, start_new_transaction=start_new_transaction, ) + if did_commit: + await self.note_transactional_commit(committable_offsets) else: await self.app.producer.flush() did_commit = await self._commit(committable_offsets) @@ -1081,6 +1083,10 @@ async def _commit_offsets( self.app.monitor.on_tp_commit(committable_offsets) return did_commit + async def note_transactional_commit(self, offsets: Mapping[TP, int]) -> None: + """Notify the transport after transactional offsets commit.""" + await self.app.consumer.note_transactional_commit(offsets) + def _filter_tps_with_pending_acks( self, topics: Optional[TPorTopicSet] = None ) -> Iterator[TP]: @@ -1286,6 +1292,9 @@ async def commit(self, tps: Mapping[TP, int]) -> bool: """Commit offsets in topic partitions.""" ... + async def note_transactional_commit(self, offsets: Mapping[TP, int]) -> None: + """Record a commit completed through the producer transaction.""" + @abc.abstractmethod async def position(self, tp: TP) -> Optional[int]: """Return the current offset for partition.""" @@ -1476,6 +1485,9 @@ async def highwaters(self, *partitions: TP) -> Mapping[TP, int]: async def _commit(self, offsets: Mapping[TP, int]) -> bool: return await self._thread.commit(offsets) + async def note_transactional_commit(self, offsets: Mapping[TP, int]) -> None: + await self._thread.note_transactional_commit(offsets) + def close(self) -> None: """Close consumer for graceful shutdown.""" self._thread.close() diff --git a/faust/transport/drivers/aiokafka.py b/faust/transport/drivers/aiokafka.py index 544dbd483..8a8aef7b1 100644 --- a/faust/transport/drivers/aiokafka.py +++ b/faust/transport/drivers/aiokafka.py @@ -791,6 +791,15 @@ async def commit(self, offsets: Mapping[TP, int]) -> bool: """Commit topic offsets.""" return await self.call_thread(self._commit, offsets) + async def note_transactional_commit(self, offsets: Mapping[TP, int]) -> None: + await self.call_thread(self._note_transactional_commit, offsets) + + def _note_transactional_commit(self, offsets: Mapping[TP, int]) -> None: + now = monotonic() + self.tp_last_committed_at.update( + {tp: now for tp in offsets if tp in self.assignment()} + ) + async def _commit(self, offsets: Mapping[TP, int]) -> bool: consumer = self._ensure_consumer() now = monotonic() diff --git a/tests/unit/transport/drivers/test_aiokafka.py b/tests/unit/transport/drivers/test_aiokafka.py index 21df98711..280d0f14e 100644 --- a/tests/unit/transport/drivers/test_aiokafka.py +++ b/tests/unit/transport/drivers/test_aiokafka.py @@ -1062,6 +1062,21 @@ async def test_commit(self, *, cthread, _consumer): with self.assert_calls_thread(cthread, _consumer, cthread._commit, offsets): await cthread.commit(offsets) + @pytest.mark.asyncio + async def test_note_transactional_commit(self, *, cthread): + offsets = {TP1: 100} + cthread.call_thread = AsyncMock() + await cthread.note_transactional_commit(offsets) + cthread.call_thread.assert_awaited_once_with( + cthread._note_transactional_commit, offsets + ) + + def test__note_transactional_commit(self, *, cthread): + cthread.assignment = Mock(return_value={TP1}) + cthread.tp_last_committed_at = {} + cthread._note_transactional_commit({TP1: 100}) + assert cthread.tp_last_committed_at[TP1] > 0 + @pytest.mark.asyncio async def test__commit(self, *, cthread, _consumer): offsets = {TP1: 1001} diff --git a/tests/unit/transport/test_consumer.py b/tests/unit/transport/test_consumer.py index d0261179a..f6a102f33 100644 --- a/tests/unit/transport/test_consumer.py +++ b/tests/unit/transport/test_consumer.py @@ -980,6 +980,9 @@ async def test_commit_offsets__did_not_commit(self, *, consumer): async def test_commit_offsets__in_transaction(self, *, consumer): consumer.in_transaction = True consumer.app.producer.flush = AsyncMock() + consumer.app.consumer = Mock( + note_transactional_commit=AsyncMock(name="note_transactional_commit") + ) consumer.transactions.commit = AsyncMock() consumer.current_assignment.update({TP1, TP2}) ret = await consumer._commit_offsets( @@ -993,6 +996,9 @@ async def test_commit_offsets__in_transaction(self, *, consumer): {TP1: 3003, TP2: 6006}, start_new_transaction=True, ) + consumer.app.consumer.note_transactional_commit.assert_awaited_once_with( + {TP1: 3003, TP2: 6006} + ) assert ret is consumer.transactions.commit.return_value @pytest.mark.asyncio