diff --git a/kafka/cluster.py b/kafka/cluster.py index ad62c2f2f..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._net.run(self._manager.wait_for, 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 4b1ec1452..9cfc8f7f9 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._net.wait_for(wakeup, timeout_ms=timeout_ms, raise_error=False) records, _ = self.fetched_records( max_records, update_offsets=update_offsets) @@ -354,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: @@ -450,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: @@ -466,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 @@ -759,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())) @@ -1039,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 992f2cb2e..b75687480 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): @@ -70,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 @@ -279,6 +281,59 @@ 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 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 + 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 + except Exception: + if raise_error: + raise + finally: + if timer is not None: + self.cancel(timer) + + 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. + """ + try: + return self.run(self.await_for, future, timeout_ms, raise_error, timeout_ms=timeout_ms) + except Exception: + if raise_error: + raise + # --- backend selection ---------------------------------------------------- diff --git a/kafka/net/manager.py b/kafka/net/manager.py index 4e3c10281..c607f75af 100644 --- a/kafka/net/manager.py +++ b/kafka/net/manager.py @@ -431,39 +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. - - 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._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) - 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 c67fb2d8c..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, 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) - coordinator._client.poll(future=future, timeout_ms=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, 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 @@ -664,21 +664,22 @@ 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) + net.wait_for(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, 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 = 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) + 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, 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, net, manager, seeded_coord, offsets): _Topic = OffsetCommitResponse.OffsetCommitResponseTopic _Partition = _Topic.OffsetCommitResponsePartition broker.respond(OffsetCommitRequest, OffsetCommitResponse( @@ -704,9 +705,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) + 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, 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 @@ -822,21 +823,22 @@ 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) + net.wait_for(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, 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 = 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) + 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, 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, net, manager, seeded_coord, partitions, offsets): _Topic = OffsetFetchResponse.OffsetFetchResponseTopic _Partition = _Topic.OffsetFetchResponsePartition _Group = OffsetFetchResponse.OffsetFetchResponseGroup @@ -874,9 +876,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) + net.wait_for(future, 5000) assert future.succeeded() assert future.value == offsets @@ -890,13 +892,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, net, 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 +927,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) + 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, 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,11 +1545,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) + net.wait_for(future, 1000, raise_error=False) 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 +1558,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 +1570,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 +1583,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 +2497,10 @@ 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) + net = coord._manager._net + net.wait_for(future, timeout_ms=5000, raise_error=False) return future @@ -2864,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() - client.poll(future=future, timeout_ms=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 fae3ae8b6..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._net.run(manager.wait_for, 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._net.run(manager.wait_for, 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._net.run(manager.wait_for, 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.run(manager.wait_for, 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._net.run(manager.wait_for, 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.run(manager.wait_for, 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.run(manager.wait_for, 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.run(manager.wait_for, 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._net.run(manager.wait_for, 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_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..089278a7e 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_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(): + fut = all_net.create_future() + all_net.call_soon(lambda: fut.success('value')) + 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.await_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 + 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 124959a26..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.run(manager.wait_for, 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.run(manager.wait_for, 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')))