fix: key pending commit work by consumer and partition

CONTEXT.md

@@ -40,10 +40,11 @@ this middleware were absent, while a skipped message is dropped mid-shutdown wit
 uncommitted for redelivery. They sound alike and their offset consequences are opposite.
 
 **Ready prefix**:
-The leading contiguous run of finished tasks on one partition — the only offsets eligible to
-commit, and all `ready` ever means in `take_ready` / `ReadyCommit` / `extract_ready_prefixes`. It
+The leading contiguous run of finished tasks on one partition for one consumer — the only offsets
+eligible to commit, and all `ready` ever means in `take_ready` / `ReadyCommit` / `extract_ready_prefixes`. It
 is not "the finished tasks": a single unfinished task at the head holds back every finished task
-behind it, which is what keeps commits from jumping past in-flight work.
+behind it on that consumer, which is what keeps commits from jumping past in-flight work. Another
+consumer's tasks on the same partition, in a different consumer group, never hold it back.
 
 **Cancellation watermark**:
 The per-`(consumer, partition)` floor recorded when a cancelled task is seen, past which offsets

README.md

@@ -113,7 +113,7 @@ This library does not install signal handlers. Shutdown is driven by your lifesp
 
 ### KafkaBatchCommitter
 
-Runs as a background asyncio task. A streaming loop absorbs `KafkaCommitTask` objects into per-partition pending state and commits each partition's contiguous-done prefix when total pending crosses `commit_batch_size`, when `commit_batch_timeout_sec` fires, or when `commit_all`/`close` sets the flush event. Cancelled tasks are a hard boundary: the offset advance stops at the cancelled task so it gets redelivered on restart (at-least-once). If the committer's task dies, `CommitterIsDeadError` is raised to callers.
+Runs as a background asyncio task. A streaming loop absorbs `KafkaCommitTask` objects into pending state per consumer and partition and commits each one's contiguous-done prefix when total pending crosses `commit_batch_size`, when `commit_batch_timeout_sec` fires, or when `commit_all`/`close` sets the flush event. Cancelled tasks are a hard boundary: the offset advance stops at the cancelled task so it gets redelivered on restart (at-least-once). If the committer's task dies, `CommitterIsDeadError` is raised to callers.
 
 ## API reference
 

faststream_concurrent_aiokafka/_pending_state.py

@@ -8,6 +8,8 @@
 
 
 _OFFSET_KEY: typing.Final = operator.attrgetter("offset")
+# (id(consumer), partition): one owner of a partition. Pending lists and watermarks share it.
+OwnerKey: typing.TypeAlias = tuple[int, TopicPartition]
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True, slots=True)
@@ -28,20 +30,20 @@ def insert_sorted(partition_pending: list[KafkaCommitTask], new_ct: KafkaCommitT
 
 
 def extract_ready_prefixes(
-    pending: dict[TopicPartition, list[KafkaCommitTask]],
-) -> tuple[dict[TopicPartition, list[KafkaCommitTask]], int]:
-    # Pending lists are maintained in offset order by insert_sorted. Per partition, find
+    pending: dict[OwnerKey, list[KafkaCommitTask]],
+) -> tuple[dict[OwnerKey, list[KafkaCommitTask]], int]:
+    # Pending lists are maintained in offset order by insert_sorted. Per owner, find
     # the first not-done task; tasks before it form the contiguous-done prefix and become
     # "ready". A cancelled task is treated as a hard boundary: cancelled + everything after
     # is dropped from pending and added to ready (so task_done() balances
     # messages_queue.join), while map_offsets_per_partition stops the offset advance at
     # the cancelled task so the uncommitted offsets get redelivered on restart
     # (at-least-once). Returns (ready, count) so the caller can update its cached
     # pending_count without re-summing list lengths.
-    ready: dict[TopicPartition, list[KafkaCommitTask]] = {}
+    ready: dict[OwnerKey, list[KafkaCommitTask]] = {}
     ready_count = 0
-    empty_partitions: list[TopicPartition] = []
-    for partition, partition_pending in pending.items():
+    empty_owners: list[OwnerKey] = []
+    for owner, partition_pending in pending.items():
         prefix_end = 0
         for index, task in enumerate(partition_pending):
             if task.asyncio_task.cancelled():
@@ -53,21 +55,21 @@ def extract_ready_prefixes(
             prefix_end = index + 1
 
         if prefix_end > 0:
-            ready[partition] = partition_pending[:prefix_end]
+            ready[owner] = partition_pending[:prefix_end]
             ready_count += prefix_end
             del partition_pending[:prefix_end]
         if not partition_pending:
-            empty_partitions.append(partition)
+            empty_owners.append(owner)
 
-    for k in empty_partitions:
-        del pending[k]
+    for owner in empty_owners:
+        del pending[owner]
     return ready, ready_count
 
 
 def map_offsets_per_partition(
     consumer_id: int,
     consumer_tasks: list[KafkaCommitTask],
-    watermarks: dict[tuple[int, TopicPartition], int],
+    watermarks: dict[OwnerKey, int],
 ) -> dict[TopicPartition, int]:
     # `watermarks` is mutated: any cancelled task seen here records (or lowers) the
     # (consumer, partition) watermark. Subsequent batches for the same consumer will see
@@ -79,7 +81,7 @@ def map_offsets_per_partition(
 
     partitions_to_offsets: dict[TopicPartition, int] = {}
     for partition, tasks in by_partition.items():
-        wm_key: tuple[int, TopicPartition] = (consumer_id, partition)
+        wm_key: OwnerKey = (consumer_id, partition)
         max_offset: int | None = None
         for task in sorted(tasks, key=_OFFSET_KEY):
             if task.asyncio_task.cancelled():
@@ -112,25 +114,25 @@ class ReadyCommit:
 
 
 class PendingCommits:
-    """Owns per-partition pending commit tasks, held failed commits, pending count, and cancellation watermarks.
+    """Owns per-owner pending commit tasks, held failed commits, pending count, and cancellation watermarks.
 
     Synchronous and single-owner: the committer's streaming loop is the sole
     mutator, so no locking is needed. Reads asyncio task state (done/cancelled)
     but never awaits and never performs I/O.
     """
 
     def __init__(self) -> None:
-        self._pending: dict[TopicPartition, list[KafkaCommitTask]] = {}
+        self._pending: dict[OwnerKey, list[KafkaCommitTask]] = {}
         self._count: int = 0
-        self._watermarks: dict[tuple[int, TopicPartition], int] = {}
+        self._watermarks: dict[OwnerKey, int] = {}
         # Commits that hit a transient error; take_ready() returns them again, merged per consumer.
         self._held: list[ReadyCommit] = []
 
     def __len__(self) -> int:
         return self._count
 
     def absorb(self, ct: KafkaCommitTask) -> None:
-        insert_sorted(self._pending.setdefault(ct.topic_partition, []), ct)
+        insert_sorted(self._pending.setdefault((id(ct.consumer), ct.topic_partition), []), ct)
         self._count += 1
 
     def hold(self, rc: ReadyCommit) -> None:

tests/test_kafka_committer.py

@@ -528,7 +528,7 @@ def test_committer_map_offsets_advances_to_max_per_partition(mock_consumer: Mock
 
 
 def test_extract_ready_prefixes_empty_pending() -> None:
-    pending: dict[TopicPartition, list[KafkaCommitTask]] = {}
+    pending: dict[_pending_state.OwnerKey, list[KafkaCommitTask]] = {}
     ready, ready_count = _pending_state.extract_ready_prefixes(pending)
     assert ready == {}
     assert ready_count == 0
@@ -537,6 +537,7 @@ def test_extract_ready_prefixes_empty_pending() -> None:
 
 def test_extract_ready_prefixes_all_done(mock_consumer: MockAIOKafkaConsumer) -> None:
     tp: typing.Final = TopicPartition(topic="t", partition=0)
+    owner: typing.Final = (id(mock_consumer), tp)
     tasks: typing.Final = [
         KafkaCommitTask(
             asyncio_task=MockAsyncioTask(done=True),  # ty: ignore[invalid-argument-type]
@@ -546,17 +547,18 @@ def test_extract_ready_prefixes_all_done(mock_consumer: MockAIOKafkaConsumer) ->
         )
         for offset in (10, 11, 12)
     ]
-    pending: dict[TopicPartition, list[KafkaCommitTask]] = {tp: list(tasks)}
+    pending: dict[_pending_state.OwnerKey, list[KafkaCommitTask]] = {owner: list(tasks)}
 
     ready, ready_count = _pending_state.extract_ready_prefixes(pending)
 
-    assert ready == {tp: tasks}
+    assert ready == {owner: tasks}
     assert ready_count == 3
     assert pending == {}  # partition emptied
 
 
 def test_extract_ready_prefixes_blocks_on_first_pending(mock_consumer: MockAIOKafkaConsumer) -> None:
     tp: typing.Final = TopicPartition(topic="t", partition=0)
+    owner: typing.Final = (id(mock_consumer), tp)
     pending_task: typing.Final = MockAsyncioTask(done=False)
     tasks: typing.Final = [
         KafkaCommitTask(
@@ -578,13 +580,13 @@ def test_extract_ready_prefixes_blocks_on_first_pending(mock_consumer: MockAIOKa
             topic_partition=tp,
         ),
     ]
-    pending: dict[TopicPartition, list[KafkaCommitTask]] = {tp: list(tasks)}
+    pending: dict[_pending_state.OwnerKey, list[KafkaCommitTask]] = {owner: list(tasks)}
 
     ready, ready_count = _pending_state.extract_ready_prefixes(pending)
 
-    assert ready == {tp: [tasks[0]]}  # only the prefix before offset 11
+    assert ready == {owner: [tasks[0]]}  # only the prefix before offset 11
     assert ready_count == 1
-    assert pending[tp] == [tasks[1], tasks[2]]
+    assert pending[owner] == [tasks[1], tasks[2]]
 
 
 def test_extract_ready_prefixes_cancelled_drops_partition(mock_consumer: MockAIOKafkaConsumer) -> None:
@@ -594,6 +596,7 @@ def test_extract_ready_prefixes_cancelled_drops_partition(mock_consumer: MockAIO
     separately stops the offset advance at the cancelled task so it gets redelivered.
     """
     tp: typing.Final = TopicPartition(topic="t", partition=0)
+    owner: typing.Final = (id(mock_consumer), tp)
     tasks: typing.Final = [
         KafkaCommitTask(
             asyncio_task=MockAsyncioTask(done=True),  # ty: ignore[invalid-argument-type]
@@ -614,11 +617,11 @@ def test_extract_ready_prefixes_cancelled_drops_partition(mock_consumer: MockAIO
             topic_partition=tp,
         ),
     ]
-    pending: dict[TopicPartition, list[KafkaCommitTask]] = {tp: list(tasks)}
+    pending: dict[_pending_state.OwnerKey, list[KafkaCommitTask]] = {owner: list(tasks)}
 
     ready, ready_count = _pending_state.extract_ready_prefixes(pending)
 
-    assert ready == {tp: tasks}  # all three included in ready
+    assert ready == {owner: tasks}  # all three included in ready
     assert ready_count == 3
     assert pending == {}  # partition emptied
 
@@ -1693,3 +1696,26 @@ async def test_a_streak_of_transient_errors_logs_one_error(caplog: pytest.LogCap
     assert all(r.exc_info is None for r in warnings)
     assert len(infos) == 1
     await committer.close()
+
+
+async def test_another_consumer_groups_in_flight_task_does_not_delay_a_commit() -> None:
+    group_a: typing.Final = MockAIOKafkaConsumer(group_id="a")
+    group_b: typing.Final = MockAIOKafkaConsumer(group_id="b")
+    committer: typing.Final = KafkaBatchCommitter(commit_batch_timeout_sec=0.05, commit_batch_size=100)
+    committer.spawn()
+    in_flight: typing.Final = asyncio.create_task(_never_finishes())
+    await committer.send_task(_commit_task(group_a, 0, in_flight))
+    await committer.send_task(
+        KafkaCommitTask(
+            asyncio_task=asyncio.create_task(_finishes()),
+            offset=101,
+            consumer=group_b,
+            topic_partition=TopicPartition(topic="t", partition=0),
+        )
+    )
+
+    await _drive_until(lambda: group_b.commit.call_count >= 1, deadline_sec=0.5)
+
+    assert group_b.commit.call_args.args[0] == {TopicPartition(topic="t", partition=0): 102}
+    group_a.commit.assert_not_called()
+    await _cancel_all(committer, in_flight)

tests/test_pending_commits.py

@@ -168,3 +168,48 @@ def test_held_commits_stay_per_consumer() -> None:
 
     assert ready[id(a)].offsets == {_tp(): 1}
     assert ready[id(b)].offsets == {_tp(): 4}
+
+
+@pytest.mark.parametrize("blocked_first", [True, False])
+def test_another_consumers_unfinished_task_does_not_hold_back_a_ready_prefix(*, blocked_first: bool) -> None:
+    blocked, ready_consumer = MockAIOKafkaConsumer(), MockAIOKafkaConsumer()
+    if not blocked_first:
+        blocked, ready_consumer = ready_consumer, blocked
+    pending = PendingCommits()
+    pending.absorb(make_commit_task(blocked, _tp(), offset=5, done=False))
+    pending.absorb(make_commit_task(ready_consumer, _tp(), offset=6, done=True))
+
+    [rc] = pending.take_ready()
+
+    assert rc.consumer is ready_consumer
+    assert rc.offsets == {_tp(): 7}
+    assert len(pending) == 1
+
+
+def test_unfinished_task_still_holds_back_its_own_consumers_later_tasks(mock_consumer: MockAIOKafkaConsumer) -> None:
+    other = MockAIOKafkaConsumer()
+    pending = PendingCommits()
+    pending.absorb(make_commit_task(mock_consumer, _tp(), offset=5, done=False))
+    pending.absorb(make_commit_task(other, _tp(), offset=6, done=True))
+    pending.absorb(make_commit_task(mock_consumer, _tp(), offset=7, done=True))
+
+    ready = {id(rc.consumer): rc for rc in pending.take_ready()}
+
+    assert set(ready) == {id(other)}
+    assert len(pending) == 2
+
+
+def test_cancelled_task_sets_a_watermark_only_for_its_own_consumer() -> None:
+    a, b = MockAIOKafkaConsumer(), MockAIOKafkaConsumer()
+    pending = PendingCommits()
+    pending.absorb(make_commit_task(a, _tp(), offset=5, cancelled=True))
+    pending.absorb(make_commit_task(b, _tp(), offset=6, done=True))
+    first = {id(rc.consumer): rc for rc in pending.take_ready()}
+
+    pending.absorb(make_commit_task(a, _tp(), offset=8, done=True))
+    pending.absorb(make_commit_task(b, _tp(), offset=9, done=True))
+    second = {id(rc.consumer): rc for rc in pending.take_ready()}
+
+    assert first[id(b)].offsets == {_tp(): 7}
+    assert second[id(a)].offsets == {}
+    assert second[id(b)].offsets == {_tp(): 10}