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}