Track transactional commits in aiokafka monitor

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()

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()

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}

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