From 78af6653e58a86cfcc4a7cb779e1ae08ee66a915 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Thu, 23 Jul 2026 17:26:35 +0000 Subject: [PATCH 1/5] net: move wait_for from manager onto the NetBackend ABC Co-Authored-By: Claude Opus 4.8 (1M context) --- kafka/net/backend/abstract.py | 41 +++++++++++++++++++++++++++ kafka/net/manager.py | 34 ++++------------------ test/net/backend/test_abstract.py | 6 ++++ test/net/backend/test_all_backends.py | 20 +++++++++++++ 4 files changed, 72 insertions(+), 29 deletions(-) diff --git a/kafka/net/backend/abstract.py b/kafka/net/backend/abstract.py index 992f2cb2e..8aed93ff9 100644 --- a/kafka/net/backend/abstract.py +++ b/kafka/net/backend/abstract.py @@ -51,6 +51,8 @@ import importlib from typing import Any, Callable, Optional, Protocol, Sequence, Tuple, runtime_checkable +import kafka.errors as Errors + @runtime_checkable class NetBackendFuture(Protocol): @@ -279,6 +281,45 @@ def create_future(self) -> NetBackendFuture: def wakeup(self) -> None: """Interrupt the loop's select() from another thread.""" + # --- shared helpers (composed from the primitives above) -------------- + async def wait_for(self, future: NetBackendFuture, timeout_ms: Optional[float]) -> Any: + """Await ``future`` with a timeout in ms. Raises KafkaTimeoutError on timeout. + + Composed entirely from ``create_future`` / ``call_later`` / ``cancel``, + so the implementation is identical across backends; it lives here rather + than in each backend (or on the manager, which merely delegates). A + backend with a native bounded-await may override. + + Must be awaited from a coroutine running on this loop. The underlying + future is not cancelled on timeout -- it continues to run; the timeout + only unblocks the awaiter. + """ + # Always await a backend-native wrapper, never ``future`` directly: + # ``future`` may be a plain thread-safe Future which isn't awaitable on + # every backend (e.g. asyncio rejects a bare ``yield self``). We touch it + # only via callbacks. (create_future() gives the backend's awaitable.) + wrapper = self.create_future() + def _on_success(value): + if not wrapper.is_done: + wrapper.success(value) + def _on_failure(exc): + if not wrapper.is_done: + wrapper.failure(exc) + future.add_callback(_on_success) + future.add_errback(_on_failure) + timer = None + if timeout_ms is not None: + def _on_timeout(): + if not wrapper.is_done: + wrapper.failure(Errors.KafkaTimeoutError( + 'Timed out after %s ms' % timeout_ms)) + timer = self.call_later(timeout_ms / 1000, _on_timeout) + try: + return await wrapper + finally: + if timer is not None: + self.cancel(timer) + # --- backend selection ---------------------------------------------------- diff --git a/kafka/net/manager.py b/kafka/net/manager.py index 4e3c10281..0344a41c4 100644 --- a/kafka/net/manager.py +++ b/kafka/net/manager.py @@ -432,37 +432,13 @@ def close(self, node_id=None, timeout_ms=None): self._net.close() async def wait_for(self, future, timeout_ms): - """Await `future` with a timeout in ms. Raises KafkaTimeoutError on timeout. + """Await ``future`` with a timeout in ms; raises KafkaTimeoutError on timeout. - Must be awaited from a coroutine running on this loop. The underlying - future is not cancelled on timeout - it continues to run; the timeout - only unblocks the awaiter. + Thin delegate to the backend's shared ``wait_for`` (see + ``NetBackend.wait_for``), mirroring the ``call_soon`` / ``create_future`` + shims. Must be awaited from a coroutine running on this loop. """ - # Always await a backend-native wrapper, never `future` directly: - # `future` may be a plain thread-safe Future which isn't awaitable on - # every backend (e.g. asyncio rejects a bare `yield self`). We touch it - # only via callbacks. (create_future() gives the backend's awaitable.) - wrapper = self._net.create_future() - def _on_success(value): - if not wrapper.is_done: - wrapper.success(value) - def _on_failure(exc): - if not wrapper.is_done: - wrapper.failure(exc) - future.add_callback(_on_success) - future.add_errback(_on_failure) - timer = None - if timeout_ms is not None: - def _on_timeout(): - if not wrapper.is_done: - wrapper.failure(Errors.KafkaTimeoutError( - 'Timed out after %s ms' % timeout_ms)) - timer = self._net.call_later(timeout_ms / 1000, _on_timeout) - try: - return await wrapper - finally: - if timer is not None: - self._net.cancel(timer) + return await self._net.wait_for(future, timeout_ms) def create_future(self): """Create a Future suitable for awaiting on the underlying loop. diff --git a/test/net/backend/test_abstract.py b/test/net/backend/test_abstract.py index 668a85590..5b177da3e 100644 --- a/test/net/backend/test_abstract.py +++ b/test/net/backend/test_abstract.py @@ -77,6 +77,12 @@ def test_call_soon_threadsafe_folded_into_call_soon(self): net = NetworkSelector() assert not hasattr(net, 'call_soon_threadsafe') + def test_wait_for_is_shared_concrete_helper(self): + # wait_for composes create_future/call_later/cancel, so it's a concrete + # method on the ABC (NOT abstract) that every backend inherits unchanged. + assert 'wait_for' not in NetBackend.__abstractmethods__ + assert NetworkSelector.wait_for is NetBackend.wait_for + class TestNetTransportContract: def test_kafkatcptransport_satisfies_transport(self): diff --git a/test/net/backend/test_all_backends.py b/test/net/backend/test_all_backends.py index fd8d0df5b..9dc9a75a3 100644 --- a/test/net/backend/test_all_backends.py +++ b/test/net/backend/test_all_backends.py @@ -11,6 +11,9 @@ runs background heartbeat / metadata-refresh coroutines that fire unscripted requests, which MockCluster answers and a scripted MockBroker would not. """ +import pytest + +import kafka.errors as Errors from kafka.consumer.subscription_state import SubscriptionState from kafka.coordinator.consumer import ConsumerCoordinator from kafka.net.compat import KafkaNetClient @@ -73,6 +76,23 @@ def test_metadata_refresh(all_net): manager.close() +def test_wait_for_success_and_timeout(all_net): + """The shared ``NetBackend.wait_for`` helper behaves identically on both + backends: a resolved future passes its value through, and a future that + never resolves raises KafkaTimeoutError once the bound elapses.""" + async def resolves(): + fut = all_net.create_future() + all_net.call_soon(lambda: fut.success('value')) + return await all_net.wait_for(fut, timeout_ms=5000) + assert all_net.run(resolves) == 'value' + + async def never_resolves(): + fut = all_net.create_future() + return await all_net.wait_for(fut, timeout_ms=20) + with pytest.raises(Errors.KafkaTimeoutError): + all_net.run(never_resolves) + + def test_consumer_group_join_assigns_partitions(all_net, metrics): """A consumer joins a MockCluster group, is elected leader, runs the assignor, and is assigned every partition -- exercising find-coordinator + From 083b0c23ef186d81b5357e8513aa85e90ab29af2 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Sat, 11 Jul 2026 15:12:26 -0700 Subject: [PATCH 2/5] manager.wait_for_blocking --- kafka/cluster.py | 2 +- kafka/consumer/fetcher.py | 6 +- kafka/net/manager.py | 7 +++ test/consumer/test_coordinator.py | 71 ++++++++++++----------- test/consumer/test_fetcher_mock_broker.py | 18 +++--- test/producer/test_sender.py | 4 +- 6 files changed, 56 insertions(+), 52 deletions(-) diff --git a/kafka/cluster.py b/kafka/cluster.py index ad62c2f2f..a74e8d862 100644 --- a/kafka/cluster.py +++ b/kafka/cluster.py @@ -437,7 +437,7 @@ def request_update(self): ``manager.wait_for(future, timeout_ms)``, which resolves it through the backend's own awaitable: on-loop: await self._manager.wait_for(cluster.request_update(), t) - off-loop: self._net.run(self._manager.wait_for, cluster.request_update(), None) + off-loop: self._manager.wait_for_blocking(cluster.request_update(), None) Many callers want only the flag+wake side effect and discard the token. On-loop callers that simply want to await a refresh can instead use the diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py index 4b1ec1452..427ced028 100644 --- a/kafka/consumer/fetcher.py +++ b/kafka/consumer/fetcher.py @@ -257,11 +257,7 @@ def _wake(_): wakeup.success(None) for fut in waited_on: fut.add_both(_wake) - - try: - self._net.run(self._manager.wait_for, wakeup, timeout_ms, timeout_ms=timeout_ms) - except Errors.KafkaTimeoutError: - pass + self._manager.wait_for_blocking(wakeup, timeout_ms) records, _ = self.fetched_records( max_records, update_offsets=update_offsets) diff --git a/kafka/net/manager.py b/kafka/net/manager.py index 0344a41c4..3e8c249ae 100644 --- a/kafka/net/manager.py +++ b/kafka/net/manager.py @@ -440,6 +440,13 @@ async def wait_for(self, future, timeout_ms): """ return await self._net.wait_for(future, timeout_ms) + def wait_for_blocking(self, future, timeout_ms): + try: + self.run(self.wait_for, future, timeout_ms, timeout_ms=timeout_ms) + except Exception: + pass + return future + def create_future(self): """Create a Future suitable for awaiting on the underlying loop. diff --git a/test/consumer/test_coordinator.py b/test/consumer/test_coordinator.py index c67fb2d8c..62b3eb8f1 100644 --- a/test/consumer/test_coordinator.py +++ b/test/consumer/test_coordinator.py @@ -502,7 +502,7 @@ def offsets(): } -def test_commit_offsets_async(mocker, coordinator, offsets): +def test_commit_offsets_async(mocker, manager, coordinator, offsets): mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False) mocker.patch.object(coordinator, 'ensure_coordinator_ready') # _send_offset_commit_request is an async coroutine; scheduled via @@ -512,7 +512,7 @@ async def fake_send(_offsets): mocker.patch.object(coordinator, '_send_offset_commit_request', side_effect=fake_send) future = coordinator.commit_offsets_async(offsets) - coordinator._client.poll(future=future, timeout_ms=1000) + manager.wait_for_blocking(future, 1000) assert coordinator._send_offset_commit_request.call_count == 1 @@ -649,7 +649,7 @@ def test_send_offset_commit_request_fail(coordinator, offsets): ((2, 0), 4), ((2, 1), 6), ], indirect=['broker']) -def test_send_offset_commit_request_versions(broker, seeded_coord, offsets, version): +def test_send_offset_commit_request_versions(broker, manager, seeded_coord, offsets, version): captured = {} _Topic = OffsetCommitResponse.OffsetCommitResponseTopic _Partition = _Topic.OffsetCommitResponsePartition @@ -664,21 +664,21 @@ def handler(api_key, api_version, correlation_id, request_bytes): ])]) broker.respond_fn(OffsetCommitRequest, handler) - future = seeded_coord._manager.call_soon( + future = manager.call_soon( seeded_coord._send_offset_commit_request, offsets) - seeded_coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.succeeded() assert captured['api_version'] == version -def test_send_offset_commit_request_failure(mocker, broker, seeded_coord, offsets): +def test_send_offset_commit_request_failure(mocker, broker, manager, seeded_coord, offsets): spy = mocker.spy(seeded_coord, '_failed_request') error = Errors.KafkaConnectionError('simulated transport failure') broker.fail_next(OffsetCommitRequest, error=error) - future = seeded_coord._manager.call_soon( + future = manager.call_soon( seeded_coord._send_offset_commit_request, offsets) - seeded_coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.failed() assert future.exception is error @@ -693,7 +693,7 @@ def test_send_offset_commit_request_failure(mocker, broker, seeded_coord, offset assert call_error is error -def test_send_offset_commit_request_success(mocker, broker, seeded_coord, offsets): +def test_send_offset_commit_request_success(mocker, broker, manager, seeded_coord, offsets): _Topic = OffsetCommitResponse.OffsetCommitResponseTopic _Partition = _Topic.OffsetCommitResponsePartition broker.respond(OffsetCommitRequest, OffsetCommitResponse( @@ -704,9 +704,9 @@ def test_send_offset_commit_request_success(mocker, broker, seeded_coord, offset ])])) spy = mocker.spy(seeded_coord, '_handle_offset_commit_response') - future = seeded_coord._manager.call_soon( + future = manager.call_soon( seeded_coord._send_offset_commit_request, offsets) - seeded_coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.succeeded() assert spy.call_count == 1 @@ -794,7 +794,7 @@ def test_send_offset_fetch_request_fail(coordinator, partitions): ((2, 5), 7), ((3, 0), 8), ], indirect=['broker']) -def test_send_offset_fetch_request_versions(broker, seeded_coord, partitions, version): +def test_send_offset_fetch_request_versions(broker, manager, seeded_coord, partitions, version): captured = {} _Topic = OffsetFetchResponse.OffsetFetchResponseTopic _Partition = _Topic.OffsetFetchResponsePartition @@ -822,21 +822,21 @@ def handler(api_key, api_version, correlation_id, request_bytes): ])])]) broker.respond_fn(OffsetFetchRequest, handler) - future = seeded_coord._manager.call_soon( + future = manager.call_soon( seeded_coord._send_offset_fetch_request, partitions) - seeded_coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.succeeded() assert captured['api_version'] == version -def test_send_offset_fetch_request_failure(mocker, broker, seeded_coord, partitions): +def test_send_offset_fetch_request_failure(mocker, broker, manager, seeded_coord, partitions): spy = mocker.spy(seeded_coord, '_failed_request') error = Errors.KafkaConnectionError('simulated transport failure') broker.fail_next(OffsetFetchRequest, error=error) - future = seeded_coord._manager.call_soon( + future = manager.call_soon( seeded_coord._send_offset_fetch_request, partitions) - seeded_coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.failed() assert future.exception is error @@ -850,7 +850,7 @@ def test_send_offset_fetch_request_failure(mocker, broker, seeded_coord, partiti assert call_error is error -def test_send_offset_fetch_request_success(mocker, broker, seeded_coord, partitions, offsets): +def test_send_offset_fetch_request_success(mocker, broker, manager, seeded_coord, partitions, offsets): _Topic = OffsetFetchResponse.OffsetFetchResponseTopic _Partition = _Topic.OffsetFetchResponsePartition _Group = OffsetFetchResponse.OffsetFetchResponseGroup @@ -874,9 +874,9 @@ def test_send_offset_fetch_request_success(mocker, broker, seeded_coord, partiti ])])])) spy = mocker.spy(seeded_coord, '_handle_offset_fetch_response') - future = seeded_coord._manager.call_soon( + future = manager.call_soon( seeded_coord._send_offset_fetch_request, partitions) - seeded_coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.succeeded() assert future.value == offsets @@ -890,13 +890,13 @@ def test_send_offset_fetch_request_success(mocker, broker, seeded_coord, partiti ('read_committed', True), ]) def test_send_offset_fetch_request_sets_require_stable( - broker, client, metrics, partitions, isolation_level, expected): + broker, manager, client, metrics, partitions, isolation_level, expected): coord = ConsumerCoordinator(client, SubscriptionState(), metrics=metrics, api_version=broker.broker_version, isolation_level=isolation_level) try: - client._manager.bootstrap(timeout_ms=5000) + manager.bootstrap(timeout_ms=5000) coord._subscription.subscribe(topics=['foobar']) coord.coordinator_id = 0 coord._generation = Generation(0, 'foobar', b'') @@ -925,9 +925,9 @@ def handler(api_key, api_version, correlation_id, request_bytes): ])])]) broker.respond_fn(OffsetFetchRequest, handler) - future = coord._manager.call_soon( + future = manager.call_soon( coord._send_offset_fetch_request, partitions) - coord._client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) assert future.succeeded() assert captured['require_stable'] is expected finally: @@ -1534,7 +1534,7 @@ async def _hang(*args, **kwargs): net.poll(timeout_ms=50) -def test_lookup_coordinator_failure(mocker, coordinator): +def test_lookup_coordinator_failure(mocker, manager, coordinator): # _send_group_coordinator_request is now an async coroutine scheduled # via manager.call_soon, so we drive the event loop to let the mock # fire before asserting on the returned future. @@ -1543,11 +1543,11 @@ async def fake_send(): mocker.patch.object(coordinator, '_send_group_coordinator_request', side_effect=fake_send) future = coordinator.lookup_coordinator() - coordinator._client.poll(future=future, timeout_ms=1000) + manager.wait_for_blocking(future, 1000) assert future.failed() -def test_do_join_and_sync_async_join_protocol_type_mismatch(request, broker, seeded_coord): +def test_do_join_and_sync_async_join_protocol_type_mismatch(request, broker, manager, seeded_coord): """KIP-559: JoinGroupResponse with mismatched protocol_type must raise InconsistentGroupProtocolError.""" request.addfinalizer(lambda: setattr(seeded_coord, 'state', MemberState.UNJOINED)) @@ -1556,10 +1556,10 @@ def test_do_join_and_sync_async_join_protocol_type_mismatch(request, broker, see protocol_type='not-consumer')) with pytest.raises(Errors.InconsistentGroupProtocolError): - seeded_coord._manager.run(seeded_coord._do_join_and_sync_async) + manager.run(seeded_coord._do_join_and_sync_async) -def test_do_join_and_sync_async_sync_protocol_type_mismatch(request, broker, seeded_coord): +def test_do_join_and_sync_async_sync_protocol_type_mismatch(request, broker, manager, seeded_coord): """KIP-559: SyncGroupResponse with mismatched protocol_type must raise.""" request.addfinalizer(lambda: setattr(seeded_coord, 'state', MemberState.UNJOINED)) broker.respond(JoinGroupRequest, _join_response_object( @@ -1568,10 +1568,10 @@ def test_do_join_and_sync_async_sync_protocol_type_mismatch(request, broker, see protocol_type='not-consumer')) with pytest.raises(Errors.InconsistentGroupProtocolError): - seeded_coord._manager.run(seeded_coord._do_join_and_sync_async) + manager.run(seeded_coord._do_join_and_sync_async) -def test_do_join_and_sync_async_sync_protocol_name_mismatch(request, broker, seeded_coord): +def test_do_join_and_sync_async_sync_protocol_name_mismatch(request, broker, manager, seeded_coord): """KIP-559: SyncGroupResponse with mismatched protocol_name must raise.""" request.addfinalizer(lambda: setattr(seeded_coord, 'state', MemberState.UNJOINED)) broker.respond(JoinGroupRequest, _join_response_object( @@ -1581,7 +1581,7 @@ def test_do_join_and_sync_async_sync_protocol_name_mismatch(request, broker, see protocol_name='roundrobin')) with pytest.raises(Errors.InconsistentGroupProtocolError): - seeded_coord._manager.run(seeded_coord._do_join_and_sync_async) + manager.run(seeded_coord._do_join_and_sync_async) # --------------------------------------------------------------------------- @@ -2495,8 +2495,9 @@ def test_close_no_autocommit_still_revokes(self, mocker, coordinator): def _dispatch_heartbeat(coord): """Dispatch a single _send_heartbeat_request and pump the network until the resulting future resolves.""" - future = coord._manager.call_soon(coord._send_heartbeat_request) - coord._client.poll(future=future, timeout_ms=5000) + manager = coord._manager + future = manager.call_soon(coord._send_heartbeat_request) + manager.wait_for_blocking(future, 5000) return future @@ -2864,7 +2865,7 @@ def test_metadata_growth_triggers_rejoin_end_to_end(net, metrics): # listener sees the change (the racing metadata update from KAFKA-3949). mock_cluster.set_metadata(topics=[_metadata_topic('t', num_partitions=2)]) future = client.cluster.request_update() - client.poll(future=future, timeout_ms=5000) + manager.wait_for_blocking(future, 5000) # The leader's assignment snapshot is now stale -> must rejoin. assert coordinator.need_rejoin() diff --git a/test/consumer/test_fetcher_mock_broker.py b/test/consumer/test_fetcher_mock_broker.py index fae3ae8b6..8c2c2731b 100644 --- a/test/consumer/test_fetcher_mock_broker.py +++ b/test/consumer/test_fetcher_mock_broker.py @@ -123,7 +123,7 @@ def test_advanced_cluster_epoch_triggers_validation_request( # the consumer's cluster cache sees the new epoch. _broker_metadata(broker, leader_epoch=5) manager.cluster.request_update() - manager._net.run(manager.wait_for, manager.cluster.request_update(), 1000) + manager.wait_for_blocking(manager.cluster.request_update(), 1000) captured = {} @@ -162,7 +162,7 @@ def test_validated_position_not_revalidated_forever( # Cluster leader epoch advances to 5; refresh the consumer's cache. _broker_metadata(broker, leader_epoch=5) - manager._net.run(manager.wait_for, manager.cluster.request_update(), 1000) + manager.wait_for_blocking(manager.cluster.request_update(), 1000) # No truncation (end_offset 100 >= position 50), and the broker # reports the requested epoch (3), NOT the current cluster epoch (5). @@ -211,7 +211,7 @@ def test_seek_forces_revalidation_of_new_position( fetcher._subscriptions.seek(tp, OffsetAndMetadata(50, '', 3)) _broker_metadata(broker, leader_epoch=5) - manager._net.run(manager.wait_for, manager.cluster.request_update(), 1000) + manager.wait_for_blocking(manager.cluster.request_update(), 1000) broker.respond(OffsetForLeaderEpochRequest, _ofle_response(error_code=0, leader_epoch=3, end_offset=100)) @@ -425,7 +425,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert 0 in requests, 'expected one fetch routed to the leader (node 0)' request, _ = requests[0] future = manager.send(request, node_id=0) - manager.run(manager.wait_for, future, 2000) + manager.wait_for_blocking(future, 2000) assert captured['api_version'] >= 11, ( 'KIP-392 requires FetchRequest v11+; got v%s' % captured.get('api_version')) @@ -452,7 +452,7 @@ def test_preferred_replica_cached_and_used_on_next_fetch( offline_replicas=[])], )]) # Re-pull metadata so the cluster cache knows about node 5. - manager._net.run(manager.wait_for, manager.cluster.request_update(), 2000) + manager.wait_for_blocking(manager.cluster.request_update(), 2000) fetcher._subscriptions.seek(tp, OffsetAndMetadata(0, '', 3)) @@ -543,7 +543,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert broker.node_id in requests request, _ = requests[broker.node_id] future = manager.send(request, node_id=broker.node_id) - manager.run(manager.wait_for, future, 2000) + manager.wait_for_blocking(future, 2000) assert captured['api_version'] >= 12, ( 'expected Fetch v12+ negotiation; got v%s' % captured.get('api_version')) @@ -569,7 +569,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): requests = fetcher._create_fetch_requests() request, _ = requests[broker.node_id] future = manager.send(request, node_id=broker.node_id) - manager.run(manager.wait_for, future, 2000) + manager.wait_for_blocking(future, 2000) assert captured['last_fetched_epoch'] == -1 @@ -596,7 +596,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): requests = fetcher._create_fetch_requests() request, _ = requests[broker.node_id] future = manager.send(request, node_id=broker.node_id) - manager.run(manager.wait_for, future, 2000) + manager.wait_for_blocking(future, 2000) assert captured['current_leader_epoch'] == -1 assert captured['last_fetched_epoch'] == 7 @@ -667,7 +667,7 @@ def test_current_leader_hint_updates_cluster_cache_on_known_broker( isr_nodes=[broker.node_id, 7], offline_replicas=[])], )]) - manager._net.run(manager.wait_for, manager.cluster.request_update(), 2000) + manager.wait_for_blocking(manager.cluster.request_update(), 2000) tp = TopicPartition(TOPIC, PARTITION) fetcher._subscriptions.seek(tp, OffsetAndMetadata(50, '', 3)) diff --git a/test/producer/test_sender.py b/test/producer/test_sender.py index 124959a26..d0618a6f4 100644 --- a/test/producer/test_sender.py +++ b/test/producer/test_sender.py @@ -164,7 +164,7 @@ def test_produce_request_negotiates_wire_version(sender, broker, manager, produc broker.respond_fn(ProduceResponse, _capture(captured)) future = manager.send(produce_request, node_id=0) - manager.run(manager.wait_for, future, 5000) + manager.wait_for_blocking(future, 5000) assert captured['api_version'] == produce_version @@ -208,7 +208,7 @@ def test_create_produce_requests_negotiates_wire_version( captured = {} broker.respond_fn(ProduceResponse, _capture(captured)) future = manager.send(request, node_id=node) - manager.run(manager.wait_for, future, 5000) + manager.wait_for_blocking(future, 5000) assert captured['api_version'] == produce_version, ( 'node %d: expected v%d got v%s' % (node, produce_version, captured.get('api_version'))) From ea01a1b38e4dcc6677e88443000ccb1df14d32d9 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Fri, 24 Jul 2026 12:17:35 -0700 Subject: [PATCH 3/5] Move manager.wait_for / wait_for_blocking -> net.await_for / net.wait_for --- kafka/cluster.py | 11 ++--- kafka/consumer/fetcher.py | 18 +++----- kafka/consumer/group.py | 18 ++------ kafka/coordinator/base.py | 9 ++-- kafka/coordinator/consumer.py | 21 ++------- kafka/net/backend/abstract.py | 29 +++++++----- kafka/net/manager.py | 16 ------- kafka/producer/kafka.py | 12 +---- test/consumer/test_coordinator.py | 43 +++++++++-------- test/consumer/test_fetcher.py | 14 +++--- test/consumer/test_fetcher_mock_broker.py | 56 +++++++++++------------ test/integration/test_sasl_integration.py | 7 +-- test/net/backend/test_all_backends.py | 4 +- test/net/test_manager.py | 2 +- test/producer/test_producer.py | 2 +- test/producer/test_sender.py | 8 ++-- 16 files changed, 109 insertions(+), 161 deletions(-) diff --git a/kafka/cluster.py b/kafka/cluster.py index a74e8d862..1e95f9201 100644 --- a/kafka/cluster.py +++ b/kafka/cluster.py @@ -428,16 +428,13 @@ def request_update(self): This is a cross-thread trigger, not a coroutine: it flags metadata as stale (changing the reported ttl()), wakes the refresh loop, and returns a token Future that resolves when the next update lands. It is safe to - call from any thread -- including user threads off the IO loop -- which - is precisely why the returned Future is a plain thread-safe handoff and - NOT a backend awaitable: a loop-affine future (create_future()) can't be - minted off the loop thread. + call from any thread -- including user threads off the IO loop. Do not ``await`` the returned Future directly. Await it at the edge via - ``manager.wait_for(future, timeout_ms)``, which resolves it through the + ``net.await_for(future, timeout_ms)``, which resolves it through the backend's own awaitable: - on-loop: await self._manager.wait_for(cluster.request_update(), t) - off-loop: self._manager.wait_for_blocking(cluster.request_update(), None) + on-loop: await self._net.await_for(cluster.request_update(), t) + off-loop: self._net.wait_for(cluster.request_update(), None) Many callers want only the flag+wake side effect and discard the token. On-loop callers that simply want to await a refresh can instead use the diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py index 427ced028..9cfc8f7f9 100644 --- a/kafka/consumer/fetcher.py +++ b/kafka/consumer/fetcher.py @@ -257,7 +257,7 @@ def _wake(_): wakeup.success(None) for fut in waited_on: fut.add_both(_wake) - self._manager.wait_for_blocking(wakeup, timeout_ms) + self._net.wait_for(wakeup, timeout_ms=timeout_ms, raise_error=False) records, _ = self.fetched_records( max_records, update_offsets=update_offsets) @@ -350,7 +350,7 @@ def reset_offsets_if_needed(self, timeout_ms=None): Returns the cached Future for the in-flight reset task (shared across concurrent callers) or None if no reset is needed. Callers may discard the Future (fire-and-forget, e.g. consumer.poll) or - await it via ``manager.wait_for(future, timeout_ms)`` to block + await it via ``net.await_for(future, timeout_ms)`` to block until resets complete (e.g. consumer.position). Arguments: @@ -446,7 +446,7 @@ async def _fetch_offsets_by_times_async(self, timestamps, timeout_ms=None): try: refresh_future = None backoff = False - offsets, retry = await self._manager.wait_for(future, timer.timeout_ms) + offsets, retry = await self._net.await_for(future, timeout_ms=timer.timeout_ms) except Errors.InvalidMetadataError: refresh_future = self._manager.cluster.request_update() except Errors.RetriableError: @@ -462,7 +462,7 @@ async def _fetch_offsets_by_times_async(self, timestamps, timeout_ms=None): if refresh_future: try: - await self._manager.wait_for(refresh_future, timer.timeout_ms) + await self._net.await_for(refresh_future, timeout_ms=timer.timeout_ms) except Errors.RetriableError: backoff = True @@ -755,10 +755,7 @@ async def _reset_offsets_async(self, timeout_ms=None): wait_ms = self.config['request_timeout_ms'] if timer.timeout_ms is not None: wait_ms = min(wait_ms, timer.timeout_ms) - try: - await self._manager.wait_for(metadata_update, wait_ms) - except Errors.KafkaTimeoutError: - pass + await self._net.await_for(metadata_update, timeout_ms=wait_ms, raise_error=False) continue log.debug('Resetting offsets for %s', set(offset_resets.keys())) @@ -1035,10 +1032,7 @@ async def _validate_offsets_async(self, timeout_ms=None): wait_ms = self.config['request_timeout_ms'] if timer.timeout_ms is not None: wait_ms = min(wait_ms, timer.timeout_ms) - try: - await self._manager.wait_for(metadata_update, wait_ms) - except Errors.KafkaTimeoutError: - pass + await self._net.await_for(metadata_update, timeout_ms=wait_ms, raise_error=False) continue log.debug('Validating offsets for %s', set(positions.keys())) diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py index bd12953c2..00000e6f8 100644 --- a/kafka/consumer/group.py +++ b/kafka/consumer/group.py @@ -786,13 +786,11 @@ def _fetch_all_topic_metadata(self, timeout_ms=None): timeout_ms = self.config['default_api_timeout_ms'] timer = Timer(timeout_ms) if self._cluster.metadata_refresh_in_progress: - future = self._cluster.request_update() - self._net.run(self._manager.wait_for, future, timer.timeout_ms, timeout_ms=timer.timeout_ms) + self._net.wait_for(self._cluster.request_update(), timeout_ms=timer.timeout_ms) stash = self._cluster.need_all_topic_metadata try: self._cluster.need_all_topic_metadata = True - future = self._cluster.request_update() - self._net.run(self._manager.wait_for, future, timer.timeout_ms, timeout_ms=timer.timeout_ms) + self._net.wait_for(self._cluster.request_update(), timeout_ms=timer.timeout_ms) finally: self._cluster.need_all_topic_metadata = stash @@ -966,10 +964,7 @@ def position(self, partition, timeout_ms=None): # past the user's deadline. reset_task = self._fetcher.reset_offsets_if_needed(timeout_ms=timer.timeout_ms) if reset_task is not None and not timer.expired: - try: - self._net.run(self._manager.wait_for, reset_task, timer.timeout_ms) - except Errors.KafkaTimeoutError: - pass + self._net.wait_for(reset_task, timeout_ms=timer.timeout_ms, raise_error=False) # Phase 3 (KIP-320): mark any positions whose cluster leader epoch # has advanced beyond the position's epoch and await the validation # RPC. Surfaces LogTruncationError to the caller if truncation is @@ -978,10 +973,7 @@ def position(self, partition, timeout_ms=None): validation_task = self._fetcher.validate_offsets_if_needed( timeout_ms=timer.timeout_ms) if validation_task is not None and not timer.expired: - try: - self._net.run(self._manager.wait_for, validation_task, timer.timeout_ms) - except Errors.KafkaTimeoutError: - pass + self._net.wait_for(validation_task, timeout_ms=timer.timeout_ms, raise_error=False) position = self._subscription.assignment[partition].position if position is not None: return position.offset @@ -1361,7 +1353,7 @@ def _refresh_committed_offsets(self, timeout_ms=None): Callers that also want the reset to complete should follow up with ``self._fetcher.reset_offsets_if_needed()`` and either await the - returned Task (e.g. via ``manager.wait_for``) or fire-and-forget. + returned Task (e.g. via ``net.await_for``) or fire-and-forget. Arguments: timeout_ms (int, optional): Milliseconds to block refreshing diff --git a/kafka/coordinator/base.py b/kafka/coordinator/base.py index 746f2ecad..d19058370 100644 --- a/kafka/coordinator/base.py +++ b/kafka/coordinator/base.py @@ -365,14 +365,14 @@ async def ensure_coordinator_ready_async(self, timeout_ms=None): future = self.lookup_coordinator() try: - await self._manager.wait_for(future, timer.timeout_ms) + await self._net.await_for(future, timeout_ms=timer.timeout_ms) except Errors.KafkaTimeoutError: return False except Errors.InvalidMetadataError as exc: log.debug('Requesting metadata for group coordinator request: %s', exc) metadata_update = self._cluster.request_update() try: - await self._manager.wait_for(metadata_update, timer.timeout_ms) + await self._net.await_for(metadata_update, timeout_ms=timer.timeout_ms) except Errors.KafkaTimeoutError: return False except Errors.RetriableError: @@ -514,8 +514,7 @@ async def join_group_async(self, timeout_ms=None): self._join_task = self._manager.call_soon(self._do_join_and_sync_async) try: - assignment_bytes = await self._manager.wait_for( - self._join_task, timer.timeout_ms) + assignment_bytes = await self._net.await_for(self._join_task, timeout_ms=timer.timeout_ms) except Errors.KafkaTimeoutError: # Timer expired; leave self._join_task in flight so the next # poll re-awaits it instead of sending a duplicate JoinGroup. @@ -1087,7 +1086,7 @@ async def maybe_leave_group_async(self, reason=None, timeout_ms=None): log.debug('Sending LeaveGroupRequest to %s: %s', self.coordinator_id, request) future = self._manager.send(request, node_id=self.coordinator_id) try: - response = await self._manager.wait_for(future, timeout_ms) + response = await self._net.await_for(future, timeout_ms=timeout_ms) self._handle_leave_group_response(response) except Errors.KafkaError as exc: log.error("LeaveGroup request failed: %s", exc) diff --git a/kafka/coordinator/consumer.py b/kafka/coordinator/consumer.py index 666d8aaf2..85fe8d286 100644 --- a/kafka/coordinator/consumer.py +++ b/kafka/coordinator/consumer.py @@ -418,10 +418,8 @@ def poll(self, timeout_ms=None): # essentially be ignored. See KAFKA-3949 for the complete # description of the problem. if self._subscription.subscribed_pattern: - metadata_update = self._cluster.request_update() try: - self._net.run( - self._manager.wait_for, metadata_update, timer.timeout_ms) + self._net.wait_for(self._cluster.request_update(), timeout_ms=timer.timeout_ms) except Errors.KafkaTimeoutError: log.debug('coordinator.poll: timeout updating metadata; returning early') return False @@ -675,14 +673,7 @@ async def fetch_committed_offsets_async(self, partitions, timeout_ms=None): else: future = self._manager.call_soon(self._send_offset_fetch_request, partitions) self._offset_fetch_futures[future_key] = future - - try: - await self._manager.wait_for(future, timer.timeout_ms) - except Errors.KafkaTimeoutError: - pass - except BaseException: - # handled below via future.is_done / retriable; cleanup happens too - pass + await self._net.await_for(future, timeout_ms=timer.timeout_ms, raise_error=False) if future.is_done: if future_key in self._offset_fetch_futures: @@ -860,13 +851,7 @@ async def _commit_offsets_sync_async(self, offsets, timeout_ms=None): await self.ensure_coordinator_ready_async(timeout_ms=timer.timeout_ms) future = self._manager.call_soon(self._send_offset_commit_request, offsets) - try: - await self._manager.wait_for(future, timer.timeout_ms) - except Errors.KafkaTimeoutError: - pass - except BaseException: - # handled below via future.is_done / retriable - pass + await self._net.await_for(future, timeout_ms=timer.timeout_ms, raise_error=False) if future.is_done: if future.succeeded(): diff --git a/kafka/net/backend/abstract.py b/kafka/net/backend/abstract.py index 8aed93ff9..be05c76cc 100644 --- a/kafka/net/backend/abstract.py +++ b/kafka/net/backend/abstract.py @@ -72,10 +72,10 @@ class NetBackendFuture(Protocol): 1. **Resolution thread.** A future from ``create_future()`` is created and resolved (``success`` / ``failure``) on the loop/IO thread only. Cross-thread handoffs (a user thread blocking on a loop result) use a - plain thread-safe ``Future`` bridged via ``manager.wait_for`` / - ``manager.run`` -- never a backend future awaited directly. Backends - whose native awaitable is loop-affine (``asyncio.Future``, Twisted - ``Deferred``) depend on this; their ``__await__`` adapter may assert it. + plain thread-safe ``Future`` bridged via ``net.wait_for`` -- never a + backend future awaited directly. Backends whose native awaitable is + loop-affine (``asyncio.Future``, Twisted ``Deferred``) depend on this; + their ``__await__`` adapter may assert it. 2. **Fan-out.** Multiple coroutines may ``await`` the same future and multiple callbacks may be registered; all are resumed / invoked. (A bare @@ -282,13 +282,8 @@ def wakeup(self) -> None: """Interrupt the loop's select() from another thread.""" # --- shared helpers (composed from the primitives above) -------------- - async def wait_for(self, future: NetBackendFuture, timeout_ms: Optional[float]) -> Any: - """Await ``future`` with a timeout in ms. Raises KafkaTimeoutError on timeout. - - Composed entirely from ``create_future`` / ``call_later`` / ``cancel``, - so the implementation is identical across backends; it lives here rather - than in each backend (or on the manager, which merely delegates). A - backend with a native bounded-await may override. + async def await_for(self, future: Any, timeout_ms: Optional[float], raise_error: bool = True) -> Any: + """Await ``future`` with a timeout in ms. Must be awaited from a coroutine running on this loop. The underlying future is not cancelled on timeout -- it continues to run; the timeout @@ -316,10 +311,22 @@ def _on_timeout(): timer = self.call_later(timeout_ms / 1000, _on_timeout) try: return await wrapper + except Exception: + if raise_error: + raise finally: if timer is not None: self.cancel(timer) + def wait_for(self, future: Any, timeout_ms: float, raise_error: bool=True) -> NetBackendFuture: + """Block until ``future`` resolves with a timeout in ms. + + Must be awaited from a coroutine running on this loop. The underlying + future is not cancelled on timeout -- it continues to run; the timeout + only unblocks the awaiter. + """ + return self.run(self.await_for, future, timeout_ms, raise_error, timeout_ms=timeout_ms) + # --- backend selection ---------------------------------------------------- diff --git a/kafka/net/manager.py b/kafka/net/manager.py index 3e8c249ae..c607f75af 100644 --- a/kafka/net/manager.py +++ b/kafka/net/manager.py @@ -431,22 +431,6 @@ def close(self, node_id=None, timeout_ms=None): if self._owns_net and not self._net.on_io_thread(): self._net.close() - async def wait_for(self, future, timeout_ms): - """Await ``future`` with a timeout in ms; raises KafkaTimeoutError on timeout. - - Thin delegate to the backend's shared ``wait_for`` (see - ``NetBackend.wait_for``), mirroring the ``call_soon`` / ``create_future`` - shims. Must be awaited from a coroutine running on this loop. - """ - return await self._net.wait_for(future, timeout_ms) - - def wait_for_blocking(self, future, timeout_ms): - try: - self.run(self.wait_for, future, timeout_ms, timeout_ms=timeout_ms) - except Exception: - pass - return future - def create_future(self): """Create a Future suitable for awaiting on the underlying loop. diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py index c6bf35c56..9a548aacd 100644 --- a/kafka/producer/kafka.py +++ b/kafka/producer/kafka.py @@ -732,11 +732,7 @@ def __getattr__(self, name): str(self), timeout) elif self._sender is not None: self._sender.initiate_close() - try: - self._manager.run(self._manager.wait_for, - self._sender._loop_future, timeout * 1000) - except Errors.KafkaTimeoutError: - pass + self._net.wait_for(self._sender._loop_future, timeout_ms=timeout * 1000, raise_error=False) if self._sender is not None and self._sender.is_running(): log.info("%s: Proceeding to force close the producer since pending" @@ -744,11 +740,7 @@ def __getattr__(self, name): str(self), timeout) self._sender.force_close() if not on_io_thread: - try: - self._manager.run(self._manager.wait_for, - self._sender._loop_future, self.config['retry_backoff_ms']) - except Errors.KafkaTimeoutError: - pass + self._net.wait_for(self._sender._loop_future, timeout_ms=self.config['retry_backoff_ms'], raise_error=False) if not on_io_thread: try: diff --git a/test/consumer/test_coordinator.py b/test/consumer/test_coordinator.py index 62b3eb8f1..1694bca94 100644 --- a/test/consumer/test_coordinator.py +++ b/test/consumer/test_coordinator.py @@ -502,7 +502,7 @@ def offsets(): } -def test_commit_offsets_async(mocker, manager, coordinator, offsets): +def test_commit_offsets_async(mocker, net, manager, coordinator, offsets): mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False) mocker.patch.object(coordinator, 'ensure_coordinator_ready') # _send_offset_commit_request is an async coroutine; scheduled via @@ -512,7 +512,7 @@ async def fake_send(_offsets): mocker.patch.object(coordinator, '_send_offset_commit_request', side_effect=fake_send) future = coordinator.commit_offsets_async(offsets) - manager.wait_for_blocking(future, 1000) + net.wait_for(future, 1000) assert coordinator._send_offset_commit_request.call_count == 1 @@ -649,7 +649,7 @@ def test_send_offset_commit_request_fail(coordinator, offsets): ((2, 0), 4), ((2, 1), 6), ], indirect=['broker']) -def test_send_offset_commit_request_versions(broker, manager, seeded_coord, offsets, version): +def test_send_offset_commit_request_versions(broker, net, manager, seeded_coord, offsets, version): captured = {} _Topic = OffsetCommitResponse.OffsetCommitResponseTopic _Partition = _Topic.OffsetCommitResponsePartition @@ -666,19 +666,20 @@ def handler(api_key, api_version, correlation_id, request_bytes): broker.respond_fn(OffsetCommitRequest, handler) future = manager.call_soon( seeded_coord._send_offset_commit_request, offsets) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert future.succeeded() assert captured['api_version'] == version -def test_send_offset_commit_request_failure(mocker, broker, manager, seeded_coord, offsets): +def test_send_offset_commit_request_failure(mocker, broker, net, manager, seeded_coord, offsets): spy = mocker.spy(seeded_coord, '_failed_request') error = Errors.KafkaConnectionError('simulated transport failure') broker.fail_next(OffsetCommitRequest, error=error) future = manager.call_soon( seeded_coord._send_offset_commit_request, offsets) - manager.wait_for_blocking(future, 5000) + with pytest.raises(Errors.KafkaConnectionError): + net.wait_for(future, 5000) assert future.failed() assert future.exception is error @@ -693,7 +694,7 @@ def test_send_offset_commit_request_failure(mocker, broker, manager, seeded_coor assert call_error is error -def test_send_offset_commit_request_success(mocker, broker, manager, seeded_coord, offsets): +def test_send_offset_commit_request_success(mocker, broker, net, manager, seeded_coord, offsets): _Topic = OffsetCommitResponse.OffsetCommitResponseTopic _Partition = _Topic.OffsetCommitResponsePartition broker.respond(OffsetCommitRequest, OffsetCommitResponse( @@ -706,7 +707,7 @@ def test_send_offset_commit_request_success(mocker, broker, manager, seeded_coor future = manager.call_soon( seeded_coord._send_offset_commit_request, offsets) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert future.succeeded() assert spy.call_count == 1 @@ -794,7 +795,7 @@ def test_send_offset_fetch_request_fail(coordinator, partitions): ((2, 5), 7), ((3, 0), 8), ], indirect=['broker']) -def test_send_offset_fetch_request_versions(broker, manager, seeded_coord, partitions, version): +def test_send_offset_fetch_request_versions(broker, net, manager, seeded_coord, partitions, version): captured = {} _Topic = OffsetFetchResponse.OffsetFetchResponseTopic _Partition = _Topic.OffsetFetchResponsePartition @@ -824,19 +825,20 @@ def handler(api_key, api_version, correlation_id, request_bytes): broker.respond_fn(OffsetFetchRequest, handler) future = manager.call_soon( seeded_coord._send_offset_fetch_request, partitions) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert future.succeeded() assert captured['api_version'] == version -def test_send_offset_fetch_request_failure(mocker, broker, manager, seeded_coord, partitions): +def test_send_offset_fetch_request_failure(mocker, broker, net, manager, seeded_coord, partitions): spy = mocker.spy(seeded_coord, '_failed_request') error = Errors.KafkaConnectionError('simulated transport failure') broker.fail_next(OffsetFetchRequest, error=error) future = manager.call_soon( seeded_coord._send_offset_fetch_request, partitions) - manager.wait_for_blocking(future, 5000) + with pytest.raises(Errors.KafkaConnectionError): + net.wait_for(future, 5000) assert future.failed() assert future.exception is error @@ -850,7 +852,7 @@ def test_send_offset_fetch_request_failure(mocker, broker, manager, seeded_coord assert call_error is error -def test_send_offset_fetch_request_success(mocker, broker, manager, seeded_coord, partitions, offsets): +def test_send_offset_fetch_request_success(mocker, broker, net, manager, seeded_coord, partitions, offsets): _Topic = OffsetFetchResponse.OffsetFetchResponseTopic _Partition = _Topic.OffsetFetchResponsePartition _Group = OffsetFetchResponse.OffsetFetchResponseGroup @@ -876,7 +878,7 @@ def test_send_offset_fetch_request_success(mocker, broker, manager, seeded_coord future = manager.call_soon( seeded_coord._send_offset_fetch_request, partitions) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert future.succeeded() assert future.value == offsets @@ -890,7 +892,7 @@ def test_send_offset_fetch_request_success(mocker, broker, manager, seeded_coord ('read_committed', True), ]) def test_send_offset_fetch_request_sets_require_stable( - broker, manager, client, metrics, partitions, isolation_level, expected): + broker, net, manager, client, metrics, partitions, isolation_level, expected): coord = ConsumerCoordinator(client, SubscriptionState(), metrics=metrics, api_version=broker.broker_version, @@ -927,7 +929,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): broker.respond_fn(OffsetFetchRequest, handler) future = manager.call_soon( coord._send_offset_fetch_request, partitions) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert future.succeeded() assert captured['require_stable'] is expected finally: @@ -1534,7 +1536,7 @@ async def _hang(*args, **kwargs): net.poll(timeout_ms=50) -def test_lookup_coordinator_failure(mocker, manager, coordinator): +def test_lookup_coordinator_failure(mocker, net, manager, coordinator): # _send_group_coordinator_request is now an async coroutine scheduled # via manager.call_soon, so we drive the event loop to let the mock # fire before asserting on the returned future. @@ -1543,7 +1545,7 @@ async def fake_send(): mocker.patch.object(coordinator, '_send_group_coordinator_request', side_effect=fake_send) future = coordinator.lookup_coordinator() - manager.wait_for_blocking(future, 1000) + net.wait_for(future, 1000, raise_error=False) assert future.failed() @@ -2497,7 +2499,8 @@ def _dispatch_heartbeat(coord): the resulting future resolves.""" manager = coord._manager future = manager.call_soon(coord._send_heartbeat_request) - manager.wait_for_blocking(future, 5000) + net = coord._manager._net + net.wait_for(future, timeout_ms=5000, raise_error=False) return future @@ -2865,7 +2868,7 @@ def test_metadata_growth_triggers_rejoin_end_to_end(net, metrics): # listener sees the change (the racing metadata update from KAFKA-3949). mock_cluster.set_metadata(topics=[_metadata_topic('t', num_partitions=2)]) future = client.cluster.request_update() - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) # The leader's assignment snapshot is now stale -> must rejoin. assert coordinator.need_rejoin() diff --git a/test/consumer/test_fetcher.py b/test/consumer/test_fetcher.py index 5ebab981c..0f119e7e2 100644 --- a/test/consumer/test_fetcher.py +++ b/test/consumer/test_fetcher.py @@ -639,7 +639,7 @@ def test_clean_done_fetch_futures_only_mutates_when_driven_on_loop(fetcher): def _capture_wakeup(fetcher, mocker): """Patch net.run to capture the wakeup Future (arg after wait_for) - without blocking, mirroring net.run(manager.wait_for, wakeup, timeout).""" + without blocking, mirroring net.wait_for(wakeup, timeout).""" captured = {} def fake_run(coro, *args, timeout_ms=None): @@ -659,7 +659,7 @@ def test_fetch_records_no_stall_when_response_arrives_before_wait(fetcher, topic Still required under the IO thread, where a fetch future can resolve at any moment. - This is faithful to the real wait: ``realistic_run`` returns immediately + This is faithful to the real wait: ``realistic_wait_for`` returns immediately only if the wakeup was already resolved (the synchronous fire), and otherwise behaves like a wait that blocks and times out. Drop the synchronous fire (e.g. by filtering already-done futures out of @@ -674,10 +674,7 @@ def test_fetch_records_no_stall_when_response_arrives_before_wait(fetcher, topic outcome = {'stalled': None} - def realistic_run(coro, wakeup, wait_timeout_ms, timeout_ms=None): - # Stand-in for net.run(manager.wait_for, wakeup, wait_timeout_ms, - # timeout_ms=...). The positional wait_for timeout and the run() backstop - # kwarg carry the same value. + def realistic_wait_for(wakeup, timeout_ms=None, raise_error=False): if wakeup.is_done: outcome['stalled'] = False # Emulate the IO thread having buffered the response that @@ -687,9 +684,10 @@ def realistic_run(coro, wakeup, wait_timeout_ms, timeout_ms=None): return None # Not resolved -> a real wait would block until timeout and raise. outcome['stalled'] = True - raise Errors.KafkaTimeoutError() + if raise_error: + raise Errors.KafkaTimeoutError() - mocker.patch.object(fetcher._net, 'run', side_effect=realistic_run) + mocker.patch.object(fetcher._net, 'wait_for', side_effect=realistic_wait_for) records, idle = fetcher.fetch_records(timeout_ms=10000) diff --git a/test/consumer/test_fetcher_mock_broker.py b/test/consumer/test_fetcher_mock_broker.py index 8c2c2731b..f1d733820 100644 --- a/test/consumer/test_fetcher_mock_broker.py +++ b/test/consumer/test_fetcher_mock_broker.py @@ -111,7 +111,7 @@ class TestKIP320OffsetValidation: """End-to-end OffsetsForLeaderEpoch flow through the wire.""" def test_advanced_cluster_epoch_triggers_validation_request( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """When the metadata-cached leader_epoch advances past the consumer's position epoch, the next ``maybe_validate_positions`` marks the partition and ``_validate_offsets_async`` issues an @@ -123,7 +123,7 @@ def test_advanced_cluster_epoch_triggers_validation_request( # the consumer's cluster cache sees the new epoch. _broker_metadata(broker, leader_epoch=5) manager.cluster.request_update() - manager.wait_for_blocking(manager.cluster.request_update(), 1000) + net.wait_for(manager.cluster.request_update(), 1000) captured = {} @@ -144,7 +144,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert fetcher._subscriptions.assignment[tp].position.offset == 50 def test_validated_position_not_revalidated_forever( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """Regression for #3106: consumer stalls after one fetched batch. After a leader election the cluster epoch advances (3 -> 5), but @@ -162,7 +162,7 @@ def test_validated_position_not_revalidated_forever( # Cluster leader epoch advances to 5; refresh the consumer's cache. _broker_metadata(broker, leader_epoch=5) - manager.wait_for_blocking(manager.cluster.request_update(), 1000) + net.wait_for(manager.cluster.request_update(), 1000) # No truncation (end_offset 100 >= position 50), and the broker # reports the requested epoch (3), NOT the current cluster epoch (5). @@ -193,11 +193,11 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert state.is_fetchable() # And no further OffsetForLeaderEpoch request is issued. - manager.run(fetcher._validate_offsets_async, 1000) + net.run(fetcher._validate_offsets_async, 1000) assert ofle_requests[0] == 1 def test_seek_forces_revalidation_of_new_position( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """A seek must re-arm validation even after a prior validation. Companion to #3106: the per-partition reconciled-leader-epoch must be @@ -211,7 +211,7 @@ def test_seek_forces_revalidation_of_new_position( fetcher._subscriptions.seek(tp, OffsetAndMetadata(50, '', 3)) _broker_metadata(broker, leader_epoch=5) - manager.wait_for_blocking(manager.cluster.request_update(), 1000) + net.wait_for(manager.cluster.request_update(), 1000) broker.respond(OffsetForLeaderEpochRequest, _ofle_response(error_code=0, leader_epoch=3, end_offset=100)) @@ -230,7 +230,7 @@ def test_seek_forces_revalidation_of_new_position( 'seeked position not re-validated -> could consume past truncation' def test_diverged_seeks_to_endpoint_with_policy( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """end_offset < position.offset (valid epoch) on the wire triggers a seek to the broker-reported divergence point - preserves progress and tags position with the confirmed epoch. Mirrors Java's @@ -252,7 +252,7 @@ def test_diverged_seeks_to_endpoint_with_policy( assert pos.leader_epoch == 3 def test_diverged_raises_when_no_reset_policy( - self, broker, client, manager, metrics): + self, broker, net, client, manager, metrics): """With offset_reset_strategy=NONE, the same wire response produces LogTruncationError carrying the divergent offset and leaves the position untouched.""" @@ -282,7 +282,7 @@ def test_diverged_raises_when_no_reset_policy( assert subs.assignment[tp].position.offset == 100 def test_undefined_response_resets_with_policy( - self, broker, manager, fetcher): + self, broker, net, fetcher): """UNDEFINED end_offset/leader_epoch (broker has no record of our epoch) with a reset policy: no known seek point, so fall back to auto_offset_reset rather than silently dropping the epoch.""" @@ -293,14 +293,14 @@ def test_undefined_response_resets_with_policy( broker.respond(OffsetForLeaderEpochRequest, _ofle_response(error_code=0, leader_epoch=-1, end_offset=-1)) - manager.run(fetcher._validate_offsets_async, 1000) + net.run(fetcher._validate_offsets_async, 1000) assert fetcher._cached_log_truncation is None assert fetcher._subscriptions.assignment[tp].awaiting_reset assert not fetcher._subscriptions.assignment[tp].awaiting_validation def test_undefined_response_raises_when_no_reset_policy( - self, broker, client, manager, metrics): + self, broker, net, client, manager, metrics): """UNDEFINED response with reset policy NONE: LogTruncationError with divergent_offsets[tp] == None (no known recovery point).""" _broker_metadata(broker, leader_epoch=3) @@ -326,7 +326,7 @@ def test_undefined_response_raises_when_no_reset_policy( assert subs.assignment[tp].position.offset == 100 def test_fenced_epoch_on_fetch_marks_validation_then_succeeds( - self, broker, manager, fetcher): + self, broker, net, fetcher): """A FENCED_LEADER_EPOCH on a real Fetch response routes through ``request_position_validation``; a subsequent ``_validate_offsets_async`` issues the OffsetForLeaderEpochRequest @@ -350,7 +350,7 @@ def test_fenced_epoch_on_fetch_marks_validation_then_succeeds( # Now run the validation driver; it should send OffsetForLeaderEpoch # and clear the flag. - manager.run(fetcher._validate_offsets_async, 1000) + net.run(fetcher._validate_offsets_async, 1000) assert not fetcher._subscriptions.assignment[tp].awaiting_validation def test_validation_retries_on_fenced_epoch_response( @@ -401,7 +401,7 @@ class TestKIP392RackAwareFetching: """End-to-end: client_rack arrives on the wire and the broker's preferred_read_replica is honored on the next fetch.""" - def test_rack_id_sent_on_fetch_request(self, broker, manager, fetcher): + def test_rack_id_sent_on_fetch_request(self, broker, net, manager, fetcher): """FetchRequest carries ``rack_id`` when client_rack is configured, and negotiates to v11+ against a modern broker.""" fetcher.config['client_rack'] = 'us-east-1a' @@ -425,14 +425,14 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert 0 in requests, 'expected one fetch routed to the leader (node 0)' request, _ = requests[0] future = manager.send(request, node_id=0) - manager.wait_for_blocking(future, 2000) + net.wait_for(future, 2000) assert captured['api_version'] >= 11, ( 'KIP-392 requires FetchRequest v11+; got v%s' % captured.get('api_version')) assert captured['rack_id'] == 'us-east-1a' def test_preferred_replica_cached_and_used_on_next_fetch( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """First fetch goes to the leader; broker returns ``preferred_read_replica=N``; second fetch routes to node N.""" tp = TopicPartition(TOPIC, PARTITION) @@ -452,7 +452,7 @@ def test_preferred_replica_cached_and_used_on_next_fetch( offline_replicas=[])], )]) # Re-pull metadata so the cluster cache knows about node 5. - manager.wait_for_blocking(manager.cluster.request_update(), 2000) + net.wait_for(manager.cluster.request_update(), 2000) fetcher._subscriptions.seek(tp, OffsetAndMetadata(0, '', 3)) @@ -473,7 +473,7 @@ def test_preferred_replica_cached_and_used_on_next_fetch( assert fetcher._select_read_replica(tp) == 5 def test_preferred_replica_negative_one_means_leader( - self, broker, manager, fetcher): + self, broker, fetcher): """``preferred_read_replica == -1`` is the broker explicitly telling the client to stop using a cached follower.""" tp = TopicPartition(TOPIC, PARTITION) @@ -518,7 +518,7 @@ class TestFetchV12Epoch: """FetchRequest v12 split-epoch request encoding and tagged response handling.""" def test_negotiates_v12_and_sends_split_epoch_fields( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """current_leader_epoch comes from cluster metadata, last_fetched_epoch from the position - and they are sent distinctly on the wire.""" tp = TopicPartition(TOPIC, PARTITION) @@ -543,7 +543,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert broker.node_id in requests request, _ = requests[broker.node_id] future = manager.send(request, node_id=broker.node_id) - manager.wait_for_blocking(future, 2000) + net.wait_for(future, 2000) assert captured['api_version'] >= 12, ( 'expected Fetch v12+ negotiation; got v%s' % captured.get('api_version')) @@ -551,7 +551,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): assert captured['last_fetched_epoch'] == 3 def test_last_fetched_epoch_is_minus_one_when_position_has_no_epoch( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """A position without a known epoch (e.g. bare seek) sends -1.""" tp = TopicPartition(TOPIC, PARTITION) fetcher._subscriptions.seek(tp, OffsetAndMetadata(0, '', -1)) @@ -569,12 +569,12 @@ def handler(api_key, api_version, correlation_id, request_bytes): requests = fetcher._create_fetch_requests() request, _ = requests[broker.node_id] future = manager.send(request, node_id=broker.node_id) - manager.wait_for_blocking(future, 2000) + net.wait_for(future, 2000) assert captured['last_fetched_epoch'] == -1 def test_current_leader_epoch_minus_one_when_metadata_has_no_epoch( - self, broker, manager, fetcher): + self, broker, net, manager, fetcher): """If the cluster cache has no epoch for the partition, send -1 (not the position epoch) - we honestly don't know the current leader.""" tp = TopicPartition(TOPIC, PARTITION) @@ -596,7 +596,7 @@ def handler(api_key, api_version, correlation_id, request_bytes): requests = fetcher._create_fetch_requests() request, _ = requests[broker.node_id] future = manager.send(request, node_id=broker.node_id) - manager.wait_for_blocking(future, 2000) + net.wait_for(future, 2000) assert captured['current_leader_epoch'] == -1 assert captured['last_fetched_epoch'] == 7 @@ -624,7 +624,7 @@ def test_diverging_epoch_response_marks_partition_for_validation( assert spy.call_count >= 1 def test_diverging_epoch_with_unset_end_offset_is_ignored( - self, broker, manager, fetcher): + self, broker, fetcher): """A divergence struct with end_offset = -1 is treated as 'no divergence reported' and records are parsed normally.""" from unittest.mock import MagicMock @@ -642,7 +642,7 @@ def test_diverging_epoch_with_unset_end_offset_is_ignored( assert not fetcher._subscriptions.assignment[tp].awaiting_validation def test_current_leader_hint_updates_cluster_cache_on_known_broker( - self, broker, manager, fetcher, mocker): + self, broker, net, manager, fetcher, mocker): """A leader-change error response carrying current_leader with a newer epoch updates the cached leader id+epoch. If the new leader id is a broker we already know, no metadata refresh is needed.""" @@ -667,7 +667,7 @@ def test_current_leader_hint_updates_cluster_cache_on_known_broker( isr_nodes=[broker.node_id, 7], offline_replicas=[])], )]) - manager.wait_for_blocking(manager.cluster.request_update(), 2000) + net.wait_for(manager.cluster.request_update(), 2000) tp = TopicPartition(TOPIC, PARTITION) fetcher._subscriptions.seek(tp, OffsetAndMetadata(50, '', 3)) diff --git a/test/integration/test_sasl_integration.py b/test/integration/test_sasl_integration.py index 504b32ce1..886383ebb 100644 --- a/test/integration/test_sasl_integration.py +++ b/test/integration/test_sasl_integration.py @@ -82,11 +82,8 @@ def test_client(request, sasl_kafka): try: manager.bootstrap(timeout_ms=5000) # auto-starts the owned net - async def fetch_metadata(): - future = manager.send(MetadataRequest(topics=None, version=1), node_id=None) - return await manager.wait_for(future, 10000) - - result = manager.run(fetch_metadata) + future = manager.send(MetadataRequest(topics=None, version=1), node_id=None) + result = manager._net.wait_for(future, timeout_ms=10000, raise_error=True) assert topic_name in [t[1] for t in result.topics] finally: manager.close() # auto-closes the owned net diff --git a/test/net/backend/test_all_backends.py b/test/net/backend/test_all_backends.py index 9dc9a75a3..492dc4474 100644 --- a/test/net/backend/test_all_backends.py +++ b/test/net/backend/test_all_backends.py @@ -83,12 +83,12 @@ def test_wait_for_success_and_timeout(all_net): async def resolves(): fut = all_net.create_future() all_net.call_soon(lambda: fut.success('value')) - return await all_net.wait_for(fut, timeout_ms=5000) + return await all_net.await_for(fut, timeout_ms=5000) assert all_net.run(resolves) == 'value' async def never_resolves(): fut = all_net.create_future() - return await all_net.wait_for(fut, timeout_ms=20) + return await all_net.await_for(fut, timeout_ms=20) with pytest.raises(Errors.KafkaTimeoutError): all_net.run(never_resolves) diff --git a/test/net/test_manager.py b/test/net/test_manager.py index 88a63e7ed..b70ebe58d 100644 --- a/test/net/test_manager.py +++ b/test/net/test_manager.py @@ -537,7 +537,7 @@ async def hangs_then_times_out(): # wait_for should fail with KafkaTimeoutError, not GeneratorExit. async def waiter(): inner = manager.call_soon(hangs_then_times_out) - return await manager.wait_for(inner, timeout_ms=50) + return await manager._net.await_for(inner, timeout_ms=50, raise_error=True) with pytest.raises(Errors.KafkaTimeoutError): manager.run(waiter) diff --git a/test/producer/test_producer.py b/test/producer/test_producer.py index 9ee013c7c..1917155ec 100644 --- a/test/producer/test_producer.py +++ b/test/producer/test_producer.py @@ -108,7 +108,7 @@ def _producer_for_send_test(partitioner): the real network.""" producer = _mock_producer(partitioner=partitioner) producer._sender.initiate_close() - producer._manager.run(producer._manager.wait_for, producer._sender._loop_future, 2000) + producer._net.wait_for(producer._sender._loop_future, 2000) producer._accumulator = MagicMock() producer._sender = MagicMock() # close() now blocks on the sender's loop Future; give the mock an diff --git a/test/producer/test_sender.py b/test/producer/test_sender.py index d0618a6f4..ca4edbc7b 100644 --- a/test/producer/test_sender.py +++ b/test/producer/test_sender.py @@ -145,7 +145,7 @@ def fn(api_key, api_version, correlation_id, request_bytes): ((0, 9), 1), ((0, 8, 0), 0), ], indirect=['broker']) -def test_produce_request_negotiates_wire_version(sender, broker, manager, produce_version): +def test_produce_request_negotiates_wire_version(sender, broker, net, manager, produce_version): """``Sender._produce_request`` returns a ProduceRequest with no fixed version; the connection negotiates the wire version against the broker's api_versions table at send time. We verify by capturing the api_version @@ -164,7 +164,7 @@ def test_produce_request_negotiates_wire_version(sender, broker, manager, produc broker.respond_fn(ProduceResponse, _capture(captured)) future = manager.send(produce_request, node_id=0) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert captured['api_version'] == produce_version @@ -176,7 +176,7 @@ def test_produce_request_negotiates_wire_version(sender, broker, manager, produc ((0, 8, 0), 0), ], indirect=['broker']) def test_create_produce_requests_negotiates_wire_version( - sender, broker, manager, produce_version): + sender, broker, net, manager, produce_version): """``_create_produce_requests`` builds one ProduceRequest per node; each one negotiates independently against its broker's api_versions table. We send each through the MockBroker (all routed to the single @@ -208,7 +208,7 @@ def test_create_produce_requests_negotiates_wire_version( captured = {} broker.respond_fn(ProduceResponse, _capture(captured)) future = manager.send(request, node_id=node) - manager.wait_for_blocking(future, 5000) + net.wait_for(future, 5000) assert captured['api_version'] == produce_version, ( 'node %d: expected v%d got v%s' % (node, produce_version, captured.get('api_version'))) From 6c50b8f387990ada947828f37748aa9e40da123c Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Tue, 28 Jul 2026 08:23:39 -0700 Subject: [PATCH 4/5] Fix raise_error in net.wait_for --- kafka/net/backend/abstract.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/kafka/net/backend/abstract.py b/kafka/net/backend/abstract.py index be05c76cc..b75687480 100644 --- a/kafka/net/backend/abstract.py +++ b/kafka/net/backend/abstract.py @@ -318,14 +318,21 @@ def _on_timeout(): if timer is not None: self.cancel(timer) - def wait_for(self, future: Any, timeout_ms: float, raise_error: bool=True) -> NetBackendFuture: - """Block until ``future`` resolves with a timeout in ms. - - Must be awaited from a coroutine running on this loop. The underlying - future is not cancelled on timeout -- it continues to run; the timeout - only unblocks the awaiter. + def wait_for(self, future: Any, timeout_ms: Optional[float], raise_error: bool=True) -> Any: + """Block the calling thread until ``future`` resolves, with a timeout in ms. + + The cross-thread blocking bridge for ``await_for``: schedules the await on + the loop and blocks the caller until it resolves, then returns its value + (or raises). Must be called from a user thread, never the IO thread + (``run`` raises ``RuntimeError`` there). The underlying future is not + cancelled on timeout -- it continues to run; the timeout only unblocks + the awaiter. """ - return self.run(self.await_for, future, timeout_ms, raise_error, timeout_ms=timeout_ms) + try: + return self.run(self.await_for, future, timeout_ms, raise_error, timeout_ms=timeout_ms) + except Exception: + if raise_error: + raise # --- backend selection ---------------------------------------------------- From cb63007b04b123f7d261be06ac5d28f6090a5415 Mon Sep 17 00:00:00 2001 From: Dana Powers Date: Tue, 28 Jul 2026 08:23:51 -0700 Subject: [PATCH 5/5] fixup test_await_for naming --- test/net/backend/test_all_backends.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/net/backend/test_all_backends.py b/test/net/backend/test_all_backends.py index 492dc4474..089278a7e 100644 --- a/test/net/backend/test_all_backends.py +++ b/test/net/backend/test_all_backends.py @@ -76,8 +76,8 @@ def test_metadata_refresh(all_net): manager.close() -def test_wait_for_success_and_timeout(all_net): - """The shared ``NetBackend.wait_for`` helper behaves identically on both +def test_await_for_success_and_timeout(all_net): + """The shared ``NetBackend.await_for`` helper behaves identically on both backends: a resolved future passes its value through, and a future that never resolves raises KafkaTimeoutError once the bound elapses.""" async def resolves():