Fix #606: never commit past an in-flight (unacked) message

faust/transport/consumer.py

@@ -1118,13 +1118,38 @@ def _new_offset(self, tp: TP) -> Optional[int]:
             # Note: acked is always kept sorted.
             # find first list of consecutive numbers
             batch = next(consecutive_numbers(acked))
+
+            # Never commit past a message that has been read but not yet
+            # acked: with concurrency the message at the head of the partition
+            # can still be processing while later offsets are already acked, so
+            # the acked run may jump right over it.  Committing then loses that
+            # in-flight message on the next restart.  Cap the batch at the
+            # smallest unacked offset for this tp.  See issue #606.
+            smallest_unacked = self._smallest_unacked_offset(tp)
+            if smallest_unacked is not None:
+                batch = [offset for offset in batch if offset < smallest_unacked]
+                if not batch:
+                    return None
+
             # remove them from the list to clean up.
             acked[: len(batch)] = []
             self._acked_index[tp].difference_update(batch)
             # return the highest commit offset
             return batch[-1] + 1
         return None
 
+    def _smallest_unacked_offset(self, tp: TP) -> Optional[int]:
+        # Smallest offset of a message that has been delivered to a stream but
+        # is not yet acknowledged, for this topic-partition (or None if there
+        # are no in-flight messages).  Used to make sure a commit never skips
+        # past an event that is still being processed.  See issue #606.
+        offsets = [
+            message.offset
+            for message in self._unacked_messages
+            if message.tp == tp and not message.acked
+        ]
+        return min(offsets) if offsets else None
+
     async def on_task_error(self, exc: BaseException) -> None:
         """Call when processing a message failed."""
         await self.commit()

tests/integration/broker/test_offset_commit.py

@@ -0,0 +1,226 @@
+"""Live-broker regression tests for offset-commit correctness (issue #606).
+
+Issue #606: :meth:`Consumer._new_offset` could return an offset *past* a
+message that was still in-flight (tracked but not yet acked), so a subsequent
+commit would persist a committed offset beyond an unprocessed message.  If the
+worker then crashed, that message was silently skipped on restart -- data loss.
+
+These tests drive faust's real aiokafka consumer against a live broker: they
+start an :class:`~faust.App`, let it get a real partition assignment, seed the
+consumer's ack bookkeeping to reproduce the in-flight scenario, run the actual
+``commit()`` path, and read the committed offset back from the broker with an
+admin client.  They assert faust never persists an offset past the smallest
+in-flight message.  Without the #606 fix the first test fails (the broker ends
+up with the too-far offset).
+
+The tests skip automatically when no Kafka broker is reachable (see the
+``broker_url`` fixture in ``conftest.py``); the dedicated CI integration job
+runs them against a real broker.
+"""
+
+import asyncio
+
+import pytest
+from aiokafka import AIOKafkaProducer, TopicPartition
+from aiokafka.admin import AIOKafkaAdminClient
+
+import faust
+from faust.types import TP
+from faust.types.tuples import Message
+
+pytestmark = [
+    pytest.mark.asyncio,
+    pytest.mark.filterwarnings("ignore::ResourceWarning"),
+]
+
+# Generous enough for a cold broker + group join in CI, short enough that a
+# genuine hang fails fast rather than sitting at the job timeout.
+ASSIGN_TIMEOUT = 30.0
+OP_TIMEOUT = 30.0
+N_MESSAGES = 5
+
+
+def _message(topic: str, tp: TP, offset: int) -> Message:
+    """Build a minimal Message for a given topic-partition/offset."""
+    return Message(
+        topic=topic,
+        partition=tp.partition,
+        offset=offset,
+        timestamp=0.0,
+        timestamp_type=0,
+        headers=[],
+        key=None,
+        value=b"x",
+        checksum=None,
+        serialized_key_size=0,
+        serialized_value_size=1,
+        tp=tp,
+    )
+
+
+async def _produce(bootstrap: str, topic: str, n: int) -> None:
+    """Produce ``n`` messages, which also auto-creates the topic."""
+    producer = AIOKafkaProducer(bootstrap_servers=bootstrap)
+    await producer.start()
+    try:
+        for i in range(n):
+            await producer.send_and_wait(topic, b"m%d" % i)
+    finally:
+        await producer.stop()
+
+
+async def _broker_committed(bootstrap: str, group: str, tp: TP):
+    """Read a group's committed offset for ``tp`` straight from the broker.
+
+    Uses the admin API so it does not join the consumer group (the faust app
+    is still a live member while this runs).  Returns ``None`` when nothing is
+    committed.
+    """
+    admin = AIOKafkaAdminClient(bootstrap_servers=bootstrap)
+    await admin.start()
+    try:
+        offsets = await admin.list_consumer_group_offsets(group)
+        meta = offsets.get(TopicPartition(tp.topic, tp.partition))
+        if meta is None or meta.offset < 0:
+            return None
+        return meta.offset
+    finally:
+        await admin.close()
+
+
+async def _start_assigned_app(broker_url: str, topic: str, group: str):
+    """Start an app whose consumer is assigned ``topic``-0, and return it.
+
+    The trivial agent gives the app a topology to subscribe to (and thus a real
+    partition assignment) without acking anything itself once the backlog is
+    drained.  Auto-commit is effectively disabled via a huge commit interval so
+    the only commit is the one the test triggers.
+    """
+    app = faust.App(
+        group,
+        broker=broker_url,
+        store="memory://",
+        value_serializer="raw",
+        topic_allow_declare=True,
+        topic_disable_leader=True,
+        web_enabled=False,
+        broker_commit_interval=9999.0,
+    )
+    source = app.topic(topic, partitions=1)
+
+    @app.agent(source)
+    async def _process(stream):
+        async for _ in stream:  # pragma: no cover - drains the backlog
+            pass
+
+    await asyncio.wait_for(app.start(), timeout=OP_TIMEOUT)
+    tp = TP(topic, 0)
+    deadline_loops = int(ASSIGN_TIMEOUT / 0.5)
+    for _ in range(deadline_loops):
+        if tp in app.consumer.assignment():
+            break
+        await asyncio.sleep(0.5)
+    assert tp in app.consumer.assignment(), "consumer never got a partition assignment"
+    # Let the agent drain the produced backlog so it does not mutate the
+    # ack bookkeeping while the test seeds its own scenario.
+    await asyncio.sleep(3)
+    return app
+
+
+async def _stop_app(app) -> None:
+    """Stop the app, clearing in-flight tracking first.
+
+    ``wait_empty()`` blocks shutdown until ``_unacked_messages`` drains; a test
+    that deliberately leaves a message in-flight must clear it or the app would
+    hang on stop.
+    """
+    consumer = app.consumer
+    consumer._unacked_messages = type(consumer._unacked_messages)()
+    await asyncio.wait_for(app.stop(), timeout=OP_TIMEOUT)
+
+
+async def test_does_not_commit_past_inflight_message(
+    broker_url, kafka_bootstrap, unique_topic, unique_group
+):
+    """#606: a commit must not advance past an in-flight (unacked) message.
+
+    Scenario: offset 2 is still being processed (tracked, unacked) while later
+    offsets 3 and 4 have been acked out of order, and the last committed offset
+    is 2.  The bug computed a new offset of 5 and committed it, skipping 2.
+    The fix withholds the commit, so the broker's committed offset must stay at
+    or below 2.
+    """
+    await _produce(kafka_bootstrap, unique_topic, N_MESSAGES)
+    app = await _start_assigned_app(broker_url, unique_topic, unique_group)
+    tp = TP(unique_topic, 0)
+    try:
+        consumer = app.consumer
+        # last committed offset is 2 (offsets 0..1 already done).
+        consumer._committed_offset[tp] = 2
+        # offsets 3 and 4 acked out of order.
+        consumer._acked[tp] = [3, 4]
+        consumer._acked_index[tp] = {3, 4}
+        # offset 2 is still in-flight (tracked, not acked).  A strong reference
+        # keeps it alive in the WeakSet of unacked messages.
+        inflight = _message(unique_topic, tp, 2)
+        consumer._unacked_messages = type(consumer._unacked_messages)()
+        consumer._unacked_messages.add(inflight)
+
+        # The fix must refuse to advance past the in-flight offset.
+        assert consumer._new_offset(tp) is None
+
+        # Re-seed (the _new_offset call above mutates the acked list) and run
+        # the real commit path against the broker.
+        consumer._committed_offset[tp] = 2
+        consumer._acked[tp] = [3, 4]
+        consumer._acked_index[tp] = {3, 4}
+        inflight = _message(unique_topic, tp, 2)
+        consumer._unacked_messages = type(consumer._unacked_messages)()
+        consumer._unacked_messages.add(inflight)
+        await asyncio.wait_for(consumer.commit({tp}), timeout=OP_TIMEOUT)
+
+        committed = await asyncio.wait_for(
+            _broker_committed(kafka_bootstrap, unique_group, tp), timeout=OP_TIMEOUT
+        )
+        assert committed is None or committed <= 2, (
+            f"committed offset {committed} advanced past the in-flight message "
+            "at offset 2 (data loss)"
+        )
+    finally:
+        await _stop_app(app)
+
+
+async def test_commits_contiguous_acks(
+    broker_url, kafka_bootstrap, unique_topic, unique_group
+):
+    """Positive control: with no in-flight message, the commit does advance.
+
+    All of offsets 0..4 are acked contiguously and nothing is in-flight, so the
+    real commit path must persist offset 5 (next-to-fetch) to the broker.  This
+    proves the negative test's "did not commit" result is meaningful and not
+    just a broken commit path.
+    """
+    await _produce(kafka_bootstrap, unique_topic, N_MESSAGES)
+    app = await _start_assigned_app(broker_url, unique_topic, unique_group)
+    tp = TP(unique_topic, 0)
+    try:
+        consumer = app.consumer
+        consumer._committed_offset[tp] = 0
+        consumer._acked[tp] = [0, 1, 2, 3, 4]
+        consumer._acked_index[tp] = {0, 1, 2, 3, 4}
+        consumer._unacked_messages = type(consumer._unacked_messages)()
+
+        assert consumer._new_offset(tp) == 5
+
+        consumer._committed_offset[tp] = 0
+        consumer._acked[tp] = [0, 1, 2, 3, 4]
+        consumer._acked_index[tp] = {0, 1, 2, 3, 4}
+        did_commit = await asyncio.wait_for(consumer.commit({tp}), timeout=OP_TIMEOUT)
+        assert did_commit
+
+        committed = await asyncio.wait_for(
+            _broker_committed(kafka_bootstrap, unique_group, tp), timeout=OP_TIMEOUT
+        )
+        assert committed == 5
+    finally:
+        await _stop_app(app)

tests/unit/transport/test_consumer.py

@@ -1104,6 +1104,30 @@ def test_new_offset_with_gaps(self, tp, acked, gaps, expected_offset, *, consume
         consumer._gap[tp] = gaps
         assert consumer._new_offset(tp) == expected_offset
 
+    def test_new_offset__withheld_for_inflight_message(self, *, consumer):
+        # Regression test for #606: with concurrency the head message can be
+        # still in-flight (read but not yet acked) while later offsets ack.
+        # Committing would jump over it and lose it on the next restart, so no
+        # offset may be committed until it is acked.
+        tp = TP1
+        consumer._committed_offset[tp] = 2585
+        consumer._acked[tp] = [2586, 2587, 2588, 2589]
+        inflight = Mock(name="msg2585", tp=tp, offset=2585, acked=False)
+        consumer._unacked_messages.add(inflight)
+
+        assert consumer._new_offset(tp) is None
+
+    def test_new_offset__commits_up_to_inflight_message(self, *, consumer):
+        # Offsets acked before the in-flight message are still committable;
+        # only the in-flight offset and everything after it is withheld.
+        tp = TP1
+        consumer._committed_offset[tp] = 100
+        consumer._acked[tp] = [100, 101, 102]
+        inflight = Mock(name="msg101", tp=tp, offset=101, acked=False)
+        consumer._unacked_messages.add(inflight)
+
+        assert consumer._new_offset(tp) == 101
+
     @pytest.mark.asyncio
     async def test_on_task_error(self, *, consumer):
         consumer.commit = AsyncMock(name="commit")