From e184711e4a0ea5d11b5ed67e3b28f7efbd470617 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sun, 4 Oct 2026 11:56:52 +0300 Subject: [PATCH] fix: key pending commit work by consumer and partition Pending lists were keyed by TopicPartition alone, so subscribers in different consumer groups reading the same topic shared one list per partition, and an unfinished task in one group held back the other group's ready prefix. Key them by (id(consumer), partition), the identity cancellation watermarks already use. Closes #96 --- CONTEXT.md | 7 +-- README.md | 2 +- .../_pending_state.py | 34 +++++++------- tests/test_kafka_committer.py | 42 +++++++++++++---- tests/test_pending_commits.py | 45 +++++++++++++++++++ 5 files changed, 102 insertions(+), 28 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 449ea8e..11a1684 100644 --- a/CONTEXT.md +++ b/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 diff --git a/README.md b/README.md index d7d6943..7f3277a 100644 --- a/README.md +++ b/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 diff --git a/faststream_concurrent_aiokafka/_pending_state.py b/faststream_concurrent_aiokafka/_pending_state.py index 81fa579..8f18871 100644 --- a/faststream_concurrent_aiokafka/_pending_state.py +++ b/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,9 +30,9 @@ 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 @@ -38,10 +40,10 @@ def extract_ready_prefixes( # 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,7 +114,7 @@ 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) @@ -120,9 +122,9 @@ class PendingCommits: """ 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] = [] @@ -130,7 +132,7 @@ 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: diff --git a/tests/test_kafka_committer.py b/tests/test_kafka_committer.py index f9b9787..4e7f6b1 100644 --- a/tests/test_kafka_committer.py +++ b/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) diff --git a/tests/test_pending_commits.py b/tests/test_pending_commits.py index 089ce47..1d512c9 100644 --- a/tests/test_pending_commits.py +++ b/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}