net: Implement await_for/wait_for in net backend

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

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()))

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

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)

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():

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 ----------------------------------------------------
 

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.
 

kafka/producer/kafka.py

@@ -732,23 +732,15 @@ 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"
                      " requests could not be completed within timeout %s.",
                      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: