744 lines
24 KiB
Python
744 lines
24 KiB
Python
|
"""Future-returning APIs for coroutines."""
|
||
|
|
|
||
|
|
# Copyright (c) PyZMQ Developers.
|
||
|
|
# Distributed under the terms of the Modified BSD License.
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import warnings
|
||
|
|
from asyncio import Future
|
||
|
|
from collections import deque
|
||
|
|
from functools import partial
|
||
|
|
from itertools import chain
|
||
|
|
from typing import Any, Awaitable, Callable, NamedTuple, TypeVar, cast, overload
|
||
|
|
|
||
|
|
import zmq as _zmq
|
||
|
|
from zmq import EVENTS, POLLIN, POLLOUT
|
||
|
|
from zmq._typing import Literal
|
||
|
|
|
||
|
|
|
||
|
|
class _FutureEvent(NamedTuple):
|
||
|
|
future: Future
|
||
|
|
kind: str
|
||
|
|
kwargs: dict
|
||
|
|
msg: Any
|
||
|
|
timer: Any
|
||
|
|
|
||
|
|
|
||
|
|
# These are incomplete classes and need a Mixin for compatibility with an eventloop
|
||
|
|
# defining the following attributes:
|
||
|
|
#
|
||
|
|
# _Future
|
||
|
|
# _READ
|
||
|
|
# _WRITE
|
||
|
|
# _default_loop()
|
||
|
|
|
||
|
|
|
||
|
|
class _Async:
|
||
|
|
"""Mixin for common async logic"""
|
||
|
|
|
||
|
|
_current_loop: Any = None
|
||
|
|
_Future: type[Future]
|
||
|
|
|
||
|
|
def _get_loop(self) -> Any:
|
||
|
|
"""Get event loop
|
||
|
|
|
||
|
|
Notice if event loop has changed,
|
||
|
|
and register init_io_state on activation of a new event loop
|
||
|
|
"""
|
||
|
|
if self._current_loop is None:
|
||
|
|
self._current_loop = self._default_loop()
|
||
|
|
self._init_io_state(self._current_loop)
|
||
|
|
return self._current_loop
|
||
|
|
current_loop = self._default_loop()
|
||
|
|
if current_loop is not self._current_loop:
|
||
|
|
# warn? This means a socket is being used in multiple loops!
|
||
|
|
self._current_loop = current_loop
|
||
|
|
self._init_io_state(current_loop)
|
||
|
|
return current_loop
|
||
|
|
|
||
|
|
def _default_loop(self) -> Any:
|
||
|
|
raise NotImplementedError("Must be implemented in a subclass")
|
||
|
|
|
||
|
|
def _init_io_state(self, loop=None) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class _AsyncPoller(_Async, _zmq.Poller):
|
||
|
|
"""Poller that returns a Future on poll, instead of blocking."""
|
||
|
|
|
||
|
|
_socket_class: type[_AsyncSocket]
|
||
|
|
_READ: int
|
||
|
|
_WRITE: int
|
||
|
|
raw_sockets: list[Any]
|
||
|
|
|
||
|
|
def _watch_raw_socket(self, loop: Any, socket: Any, evt: int, f: Callable) -> None:
|
||
|
|
"""Schedule callback for a raw socket"""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def _unwatch_raw_sockets(self, loop: Any, *sockets: Any) -> None:
|
||
|
|
"""Unschedule callback for a raw socket"""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def poll(self, timeout=-1) -> Awaitable[list[tuple[Any, int]]]: # type: ignore
|
||
|
|
"""Return a Future for a poll event"""
|
||
|
|
future = self._Future()
|
||
|
|
if timeout == 0:
|
||
|
|
try:
|
||
|
|
result = super().poll(0)
|
||
|
|
except Exception as e:
|
||
|
|
future.set_exception(e)
|
||
|
|
else:
|
||
|
|
future.set_result(result)
|
||
|
|
return future
|
||
|
|
|
||
|
|
loop = self._get_loop()
|
||
|
|
|
||
|
|
# register Future to be called as soon as any event is available on any socket
|
||
|
|
watcher = self._Future()
|
||
|
|
|
||
|
|
# watch raw sockets:
|
||
|
|
raw_sockets: list[Any] = []
|
||
|
|
|
||
|
|
def wake_raw(*args):
|
||
|
|
if not watcher.done():
|
||
|
|
watcher.set_result(None)
|
||
|
|
|
||
|
|
watcher.add_done_callback(
|
||
|
|
lambda f: self._unwatch_raw_sockets(loop, *raw_sockets)
|
||
|
|
)
|
||
|
|
|
||
|
|
wrapped_sockets: list[_AsyncSocket] = []
|
||
|
|
|
||
|
|
def _clear_wrapper_io(f):
|
||
|
|
for s in wrapped_sockets:
|
||
|
|
s._clear_io_state()
|
||
|
|
|
||
|
|
for socket, mask in self.sockets:
|
||
|
|
if isinstance(socket, _zmq.Socket):
|
||
|
|
if not isinstance(socket, self._socket_class):
|
||
|
|
# it's a blocking zmq.Socket, wrap it in async
|
||
|
|
socket = self._socket_class.from_socket(socket)
|
||
|
|
wrapped_sockets.append(socket)
|
||
|
|
if mask & _zmq.POLLIN:
|
||
|
|
socket._add_recv_event('poll', future=watcher)
|
||
|
|
if mask & _zmq.POLLOUT:
|
||
|
|
socket._add_send_event('poll', future=watcher)
|
||
|
|
else:
|
||
|
|
raw_sockets.append(socket)
|
||
|
|
evt = 0
|
||
|
|
if mask & _zmq.POLLIN:
|
||
|
|
evt |= self._READ
|
||
|
|
if mask & _zmq.POLLOUT:
|
||
|
|
evt |= self._WRITE
|
||
|
|
self._watch_raw_socket(loop, socket, evt, wake_raw)
|
||
|
|
|
||
|
|
def on_poll_ready(f):
|
||
|
|
if future.done():
|
||
|
|
return
|
||
|
|
if watcher.cancelled():
|
||
|
|
try:
|
||
|
|
future.cancel()
|
||
|
|
except RuntimeError:
|
||
|
|
# RuntimeError may be called during teardown
|
||
|
|
pass
|
||
|
|
return
|
||
|
|
if watcher.exception():
|
||
|
|
future.set_exception(watcher.exception())
|
||
|
|
else:
|
||
|
|
try:
|
||
|
|
result = super(_AsyncPoller, self).poll(0)
|
||
|
|
except Exception as e:
|
||
|
|
future.set_exception(e)
|
||
|
|
else:
|
||
|
|
future.set_result(result)
|
||
|
|
|
||
|
|
watcher.add_done_callback(on_poll_ready)
|
||
|
|
|
||
|
|
if wrapped_sockets:
|
||
|
|
watcher.add_done_callback(_clear_wrapper_io)
|
||
|
|
|
||
|
|
if timeout is not None and timeout > 0:
|
||
|
|
# schedule cancel to fire on poll timeout, if any
|
||
|
|
def trigger_timeout():
|
||
|
|
if not watcher.done():
|
||
|
|
watcher.set_result(None)
|
||
|
|
|
||
|
|
timeout_handle = loop.call_later(1e-3 * timeout, trigger_timeout)
|
||
|
|
|
||
|
|
def cancel_timeout(f):
|
||
|
|
if hasattr(timeout_handle, 'cancel'):
|
||
|
|
timeout_handle.cancel()
|
||
|
|
else:
|
||
|
|
loop.remove_timeout(timeout_handle)
|
||
|
|
|
||
|
|
future.add_done_callback(cancel_timeout)
|
||
|
|
|
||
|
|
def cancel_watcher(f):
|
||
|
|
if not watcher.done():
|
||
|
|
watcher.cancel()
|
||
|
|
|
||
|
|
future.add_done_callback(cancel_watcher)
|
||
|
|
|
||
|
|
return future
|
||
|
|
|
||
|
|
|
||
|
|
class _NoTimer:
|
||
|
|
@staticmethod
|
||
|
|
def cancel():
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
T = TypeVar("T", bound="_AsyncSocket")
|
||
|
|
|
||
|
|
|
||
|
|
class _AsyncSocket(_Async, _zmq.Socket[Future]):
|
||
|
|
# Warning : these class variables are only here to allow to call super().__setattr__.
|
||
|
|
# They be overridden at instance initialization and not shared in the whole class
|
||
|
|
_recv_futures = None
|
||
|
|
_send_futures = None
|
||
|
|
_state = 0
|
||
|
|
_shadow_sock: _zmq.Socket
|
||
|
|
_poller_class = _AsyncPoller
|
||
|
|
_fd = None
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
context=None,
|
||
|
|
socket_type=-1,
|
||
|
|
io_loop=None,
|
||
|
|
_from_socket: _zmq.Socket | None = None,
|
||
|
|
**kwargs,
|
||
|
|
) -> None:
|
||
|
|
if isinstance(context, _zmq.Socket):
|
||
|
|
context, _from_socket = (None, context)
|
||
|
|
if _from_socket is not None:
|
||
|
|
super().__init__(shadow=_from_socket.underlying) # type: ignore
|
||
|
|
self._shadow_sock = _from_socket
|
||
|
|
else:
|
||
|
|
super().__init__(context, socket_type, **kwargs) # type: ignore
|
||
|
|
self._shadow_sock = _zmq.Socket.shadow(self.underlying)
|
||
|
|
|
||
|
|
if io_loop is not None:
|
||
|
|
warnings.warn(
|
||
|
|
f"{self.__class__.__name__}(io_loop) argument is deprecated in pyzmq 22.2."
|
||
|
|
" The currently active loop will always be used.",
|
||
|
|
DeprecationWarning,
|
||
|
|
stacklevel=3,
|
||
|
|
)
|
||
|
|
self._recv_futures = deque()
|
||
|
|
self._send_futures = deque()
|
||
|
|
self._state = 0
|
||
|
|
self._fd = self._shadow_sock.FD
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_socket(cls: type[T], socket: _zmq.Socket, io_loop: Any = None) -> T:
|
||
|
|
"""Create an async socket from an existing Socket"""
|
||
|
|
return cls(_from_socket=socket, io_loop=io_loop)
|
||
|
|
|
||
|
|
def close(self, linger: int | None = None) -> None:
|
||
|
|
if not self.closed and self._fd is not None:
|
||
|
|
event_list: list[_FutureEvent] = list(
|
||
|
|
chain(self._recv_futures or [], self._send_futures or [])
|
||
|
|
)
|
||
|
|
for event in event_list:
|
||
|
|
if not event.future.done():
|
||
|
|
try:
|
||
|
|
event.future.cancel()
|
||
|
|
except RuntimeError:
|
||
|
|
# RuntimeError may be called during teardown
|
||
|
|
pass
|
||
|
|
self._clear_io_state()
|
||
|
|
super().close(linger=linger)
|
||
|
|
|
||
|
|
close.__doc__ = _zmq.Socket.close.__doc__
|
||
|
|
|
||
|
|
def get(self, key):
|
||
|
|
result = super().get(key)
|
||
|
|
if key == EVENTS:
|
||
|
|
self._schedule_remaining_events(result)
|
||
|
|
return result
|
||
|
|
|
||
|
|
get.__doc__ = _zmq.Socket.get.__doc__
|
||
|
|
|
||
|
|
@overload # type: ignore
|
||
|
|
def recv_multipart(
|
||
|
|
self, flags: int = 0, *, track: bool = False
|
||
|
|
) -> Awaitable[list[bytes]]: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def recv_multipart(
|
||
|
|
self, flags: int = 0, *, copy: Literal[True], track: bool = False
|
||
|
|
) -> Awaitable[list[bytes]]: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def recv_multipart(
|
||
|
|
self, flags: int = 0, *, copy: Literal[False], track: bool = False
|
||
|
|
) -> Awaitable[list[_zmq.Frame]]: # type: ignore
|
||
|
|
...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def recv_multipart(
|
||
|
|
self, flags: int = 0, copy: bool = True, track: bool = False
|
||
|
|
) -> Awaitable[list[bytes] | list[_zmq.Frame]]: ...
|
||
|
|
|
||
|
|
def recv_multipart(
|
||
|
|
self, flags: int = 0, copy: bool = True, track: bool = False
|
||
|
|
) -> Awaitable[list[bytes] | list[_zmq.Frame]]:
|
||
|
|
"""Receive a complete multipart zmq message.
|
||
|
|
|
||
|
|
Returns a Future whose result will be a multipart message.
|
||
|
|
"""
|
||
|
|
return self._add_recv_event(
|
||
|
|
'recv_multipart', dict(flags=flags, copy=copy, track=track)
|
||
|
|
)
|
||
|
|
|
||
|
|
@overload # type: ignore
|
||
|
|
def recv(self, flags: int = 0, *, track: bool = False) -> Awaitable[bytes]: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def recv(
|
||
|
|
self, flags: int = 0, *, copy: Literal[True], track: bool = False
|
||
|
|
) -> Awaitable[bytes]: ...
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def recv(
|
||
|
|
self, flags: int = 0, *, copy: Literal[False], track: bool = False
|
||
|
|
) -> Awaitable[_zmq.Frame]: ...
|
||
|
|
|
||
|
|
def recv( # type: ignore
|
||
|
|
self, flags: int = 0, copy: bool = True, track: bool = False
|
||
|
|
) -> Awaitable[bytes | _zmq.Frame]:
|
||
|
|
"""Receive a single zmq frame.
|
||
|
|
|
||
|
|
Returns a Future, whose result will be the received frame.
|
||
|
|
|
||
|
|
Recommend using recv_multipart instead.
|
||
|
|
"""
|
||
|
|
return self._add_recv_event('recv', dict(flags=flags, copy=copy, track=track))
|
||
|
|
|
||
|
|
def send_multipart( # type: ignore
|
||
|
|
self, msg_parts: Any, flags: int = 0, copy: bool = True, track=False, **kwargs
|
||
|
|
) -> Awaitable[_zmq.MessageTracker | None]:
|
||
|
|
"""Send a complete multipart zmq message.
|
||
|
|
|
||
|
|
Returns a Future that resolves when sending is complete.
|
||
|
|
"""
|
||
|
|
kwargs['flags'] = flags
|
||
|
|
kwargs['copy'] = copy
|
||
|
|
kwargs['track'] = track
|
||
|
|
return self._add_send_event('send_multipart', msg=msg_parts, kwargs=kwargs)
|
||
|
|
|
||
|
|
def send( # type: ignore
|
||
|
|
self,
|
||
|
|
data: Any,
|
||
|
|
flags: int = 0,
|
||
|
|
copy: bool = True,
|
||
|
|
track: bool = False,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> Awaitable[_zmq.MessageTracker | None]:
|
||
|
|
"""Send a single zmq frame.
|
||
|
|
|
||
|
|
Returns a Future that resolves when sending is complete.
|
||
|
|
|
||
|
|
Recommend using send_multipart instead.
|
||
|
|
"""
|
||
|
|
kwargs['flags'] = flags
|
||
|
|
kwargs['copy'] = copy
|
||
|
|
kwargs['track'] = track
|
||
|
|
kwargs.update(dict(flags=flags, copy=copy, track=track))
|
||
|
|
return self._add_send_event('send', msg=data, kwargs=kwargs)
|
||
|
|
|
||
|
|
def _deserialize(self, recvd, load):
|
||
|
|
"""Deserialize with Futures"""
|
||
|
|
f = self._Future()
|
||
|
|
|
||
|
|
def _chain(_):
|
||
|
|
"""Chain result through serialization to recvd"""
|
||
|
|
if f.done():
|
||
|
|
# chained future may be cancelled, which means nobody is going to get this result
|
||
|
|
# if it's an error, that's no big deal (probably zmq.Again),
|
||
|
|
# but if it's a successful recv, this is a dropped message!
|
||
|
|
if not recvd.cancelled() and recvd.exception() is None:
|
||
|
|
warnings.warn(
|
||
|
|
# is there a useful stacklevel?
|
||
|
|
# ideally, it would point to where `f.cancel()` was called
|
||
|
|
f"Future {f} completed while awaiting {recvd}. A message has been dropped!",
|
||
|
|
RuntimeWarning,
|
||
|
|
)
|
||
|
|
return
|
||
|
|
if recvd.exception():
|
||
|
|
f.set_exception(recvd.exception())
|
||
|
|
else:
|
||
|
|
buf = recvd.result()
|
||
|
|
try:
|
||
|
|
loaded = load(buf)
|
||
|
|
except Exception as e:
|
||
|
|
f.set_exception(e)
|
||
|
|
else:
|
||
|
|
f.set_result(loaded)
|
||
|
|
|
||
|
|
recvd.add_done_callback(_chain)
|
||
|
|
|
||
|
|
def _chain_cancel(_):
|
||
|
|
"""Chain cancellation from f to recvd"""
|
||
|
|
if recvd.done():
|
||
|
|
return
|
||
|
|
if f.cancelled():
|
||
|
|
recvd.cancel()
|
||
|
|
|
||
|
|
f.add_done_callback(_chain_cancel)
|
||
|
|
|
||
|
|
return f
|
||
|
|
|
||
|
|
def poll(self, timeout=None, flags=_zmq.POLLIN) -> Awaitable[int]: # type: ignore
|
||
|
|
"""poll the socket for events
|
||
|
|
|
||
|
|
returns a Future for the poll results.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if self.closed:
|
||
|
|
raise _zmq.ZMQError(_zmq.ENOTSUP)
|
||
|
|
|
||
|
|
p = self._poller_class()
|
||
|
|
p.register(self, flags)
|
||
|
|
poll_future = cast(Future, p.poll(timeout))
|
||
|
|
|
||
|
|
future = self._Future()
|
||
|
|
|
||
|
|
def unwrap_result(f):
|
||
|
|
if future.done():
|
||
|
|
return
|
||
|
|
if poll_future.cancelled():
|
||
|
|
try:
|
||
|
|
future.cancel()
|
||
|
|
except RuntimeError:
|
||
|
|
# RuntimeError may be called during teardown
|
||
|
|
pass
|
||
|
|
return
|
||
|
|
if f.exception():
|
||
|
|
future.set_exception(poll_future.exception())
|
||
|
|
else:
|
||
|
|
evts = dict(poll_future.result())
|
||
|
|
future.set_result(evts.get(self, 0))
|
||
|
|
|
||
|
|
if poll_future.done():
|
||
|
|
# hook up result if already done
|
||
|
|
unwrap_result(poll_future)
|
||
|
|
else:
|
||
|
|
poll_future.add_done_callback(unwrap_result)
|
||
|
|
|
||
|
|
def cancel_poll(future):
|
||
|
|
"""Cancel underlying poll if request has been cancelled"""
|
||
|
|
if not poll_future.done():
|
||
|
|
try:
|
||
|
|
poll_future.cancel()
|
||
|
|
except RuntimeError:
|
||
|
|
# RuntimeError may be called during teardown
|
||
|
|
pass
|
||
|
|
|
||
|
|
future.add_done_callback(cancel_poll)
|
||
|
|
|
||
|
|
return future
|
||
|
|
|
||
|
|
# overrides only necessary for updated types
|
||
|
|
def recv_string(self, *args, **kwargs) -> Awaitable[str]: # type: ignore
|
||
|
|
return super().recv_string(*args, **kwargs) # type: ignore
|
||
|
|
|
||
|
|
def send_string( # type: ignore
|
||
|
|
self, s: str, flags: int = 0, encoding: str = 'utf-8'
|
||
|
|
) -> Awaitable[None]:
|
||
|
|
return super().send_string(s, flags=flags, encoding=encoding) # type: ignore
|
||
|
|
|
||
|
|
def _add_timeout(self, future, timeout):
|
||
|
|
"""Add a timeout for a send or recv Future"""
|
||
|
|
|
||
|
|
def future_timeout():
|
||
|
|
if future.done():
|
||
|
|
# future already resolved, do nothing
|
||
|
|
return
|
||
|
|
|
||
|
|
# raise EAGAIN
|
||
|
|
future.set_exception(_zmq.Again())
|
||
|
|
|
||
|
|
return self._call_later(timeout, future_timeout)
|
||
|
|
|
||
|
|
def _call_later(self, delay, callback):
|
||
|
|
"""Schedule a function to be called later
|
||
|
|
|
||
|
|
Override for different IOLoop implementations
|
||
|
|
|
||
|
|
Tornado and asyncio happen to both have ioloop.call_later
|
||
|
|
with the same signature.
|
||
|
|
"""
|
||
|
|
return self._get_loop().call_later(delay, callback)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _remove_finished_future(future, event_list, event=None):
|
||
|
|
"""Make sure that futures are removed from the event list when they resolve
|
||
|
|
|
||
|
|
Avoids delaying cleanup until the next send/recv event,
|
||
|
|
which may never come.
|
||
|
|
"""
|
||
|
|
# "future" instance is shared between sockets, but each socket has its own event list.
|
||
|
|
if not event_list:
|
||
|
|
return
|
||
|
|
# only unconsumed events (e.g. cancelled calls)
|
||
|
|
# will be present when this happens
|
||
|
|
try:
|
||
|
|
event_list.remove(event)
|
||
|
|
except ValueError:
|
||
|
|
# usually this will have been removed by being consumed
|
||
|
|
return
|
||
|
|
|
||
|
|
def _add_recv_event(self, kind, kwargs=None, future=None):
|
||
|
|
"""Add a recv event, returning the corresponding Future"""
|
||
|
|
f = future or self._Future()
|
||
|
|
if kind.startswith('recv') and kwargs.get('flags', 0) & _zmq.DONTWAIT:
|
||
|
|
# short-circuit non-blocking calls
|
||
|
|
recv = getattr(self._shadow_sock, kind)
|
||
|
|
try:
|
||
|
|
r = recv(**kwargs)
|
||
|
|
except Exception as e:
|
||
|
|
f.set_exception(e)
|
||
|
|
else:
|
||
|
|
f.set_result(r)
|
||
|
|
return f
|
||
|
|
|
||
|
|
timer = _NoTimer
|
||
|
|
if hasattr(_zmq, 'RCVTIMEO'):
|
||
|
|
timeout_ms = self._shadow_sock.rcvtimeo
|
||
|
|
if timeout_ms >= 0:
|
||
|
|
timer = self._add_timeout(f, timeout_ms * 1e-3)
|
||
|
|
|
||
|
|
# we add it to the list of futures before we add the timeout as the
|
||
|
|
# timeout will remove the future from recv_futures to avoid leaks
|
||
|
|
_future_event = _FutureEvent(f, kind, kwargs, msg=None, timer=timer)
|
||
|
|
self._recv_futures.append(_future_event)
|
||
|
|
|
||
|
|
if self._shadow_sock.get(EVENTS) & POLLIN:
|
||
|
|
# recv immediately, if we can
|
||
|
|
self._handle_recv()
|
||
|
|
if self._recv_futures and _future_event in self._recv_futures:
|
||
|
|
# Don't let the Future sit in _recv_events after it's done
|
||
|
|
# no need to register this if we've already been handled
|
||
|
|
# (i.e. immediately-resolved recv)
|
||
|
|
f.add_done_callback(
|
||
|
|
partial(
|
||
|
|
self._remove_finished_future,
|
||
|
|
event_list=self._recv_futures,
|
||
|
|
event=_future_event,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
self._add_io_state(POLLIN)
|
||
|
|
return f
|
||
|
|
|
||
|
|
def _add_send_event(self, kind, msg=None, kwargs=None, future=None):
|
||
|
|
"""Add a send event, returning the corresponding Future"""
|
||
|
|
f = future or self._Future()
|
||
|
|
# attempt send with DONTWAIT if no futures are waiting
|
||
|
|
# short-circuit for sends that will resolve immediately
|
||
|
|
# only call if no send Futures are waiting
|
||
|
|
if kind in ('send', 'send_multipart') and not self._send_futures:
|
||
|
|
flags = kwargs.get('flags', 0)
|
||
|
|
nowait_kwargs = kwargs.copy()
|
||
|
|
nowait_kwargs['flags'] = flags | _zmq.DONTWAIT
|
||
|
|
|
||
|
|
# short-circuit non-blocking calls
|
||
|
|
send = getattr(self._shadow_sock, kind)
|
||
|
|
# track if the send resolved or not
|
||
|
|
# (EAGAIN if DONTWAIT is not set should proceed with)
|
||
|
|
finish_early = True
|
||
|
|
try:
|
||
|
|
r = send(msg, **nowait_kwargs)
|
||
|
|
except _zmq.Again as e:
|
||
|
|
if flags & _zmq.DONTWAIT:
|
||
|
|
f.set_exception(e)
|
||
|
|
else:
|
||
|
|
# EAGAIN raised and DONTWAIT not requested,
|
||
|
|
# proceed with async send
|
||
|
|
finish_early = False
|
||
|
|
except Exception as e:
|
||
|
|
f.set_exception(e)
|
||
|
|
else:
|
||
|
|
f.set_result(r)
|
||
|
|
|
||
|
|
if finish_early:
|
||
|
|
# short-circuit resolved, return finished Future
|
||
|
|
# schedule wake for recv if there are any receivers waiting
|
||
|
|
if self._recv_futures:
|
||
|
|
self._schedule_remaining_events()
|
||
|
|
return f
|
||
|
|
|
||
|
|
timer = _NoTimer
|
||
|
|
if hasattr(_zmq, 'SNDTIMEO'):
|
||
|
|
timeout_ms = self._shadow_sock.get(_zmq.SNDTIMEO)
|
||
|
|
if timeout_ms >= 0:
|
||
|
|
timer = self._add_timeout(f, timeout_ms * 1e-3)
|
||
|
|
|
||
|
|
# we add it to the list of futures before we add the timeout as the
|
||
|
|
# timeout will remove the future from recv_futures to avoid leaks
|
||
|
|
_future_event = _FutureEvent(f, kind, kwargs=kwargs, msg=msg, timer=timer)
|
||
|
|
self._send_futures.append(_future_event)
|
||
|
|
# Don't let the Future sit in _send_futures after it's done
|
||
|
|
f.add_done_callback(
|
||
|
|
partial(
|
||
|
|
self._remove_finished_future,
|
||
|
|
event_list=self._send_futures,
|
||
|
|
event=_future_event,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
self._add_io_state(POLLOUT)
|
||
|
|
return f
|
||
|
|
|
||
|
|
def _handle_recv(self):
|
||
|
|
"""Handle recv events"""
|
||
|
|
if not self._shadow_sock.get(EVENTS) & POLLIN:
|
||
|
|
# event triggered, but state may have been changed between trigger and callback
|
||
|
|
return
|
||
|
|
f = None
|
||
|
|
while self._recv_futures:
|
||
|
|
f, kind, kwargs, _, timer = self._recv_futures.popleft()
|
||
|
|
# skip any cancelled futures
|
||
|
|
if f.done():
|
||
|
|
f = None
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
|
||
|
|
if not self._recv_futures:
|
||
|
|
self._drop_io_state(POLLIN)
|
||
|
|
|
||
|
|
if f is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
timer.cancel()
|
||
|
|
|
||
|
|
if kind == 'poll':
|
||
|
|
# on poll event, just signal ready, nothing else.
|
||
|
|
f.set_result(None)
|
||
|
|
return
|
||
|
|
elif kind == 'recv_multipart':
|
||
|
|
recv = self._shadow_sock.recv_multipart
|
||
|
|
elif kind == 'recv':
|
||
|
|
recv = self._shadow_sock.recv
|
||
|
|
else:
|
||
|
|
raise ValueError(f"Unhandled recv event type: {kind!r}")
|
||
|
|
|
||
|
|
kwargs['flags'] |= _zmq.DONTWAIT
|
||
|
|
try:
|
||
|
|
result = recv(**kwargs)
|
||
|
|
except Exception as e:
|
||
|
|
f.set_exception(e)
|
||
|
|
else:
|
||
|
|
f.set_result(result)
|
||
|
|
|
||
|
|
def _handle_send(self):
|
||
|
|
if not self._shadow_sock.get(EVENTS) & POLLOUT:
|
||
|
|
# event triggered, but state may have been changed between trigger and callback
|
||
|
|
return
|
||
|
|
f = None
|
||
|
|
while self._send_futures:
|
||
|
|
f, kind, kwargs, msg, timer = self._send_futures.popleft()
|
||
|
|
# skip any cancelled futures
|
||
|
|
if f.done():
|
||
|
|
f = None
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
|
||
|
|
if not self._send_futures:
|
||
|
|
self._drop_io_state(POLLOUT)
|
||
|
|
|
||
|
|
if f is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
timer.cancel()
|
||
|
|
|
||
|
|
if kind == 'poll':
|
||
|
|
# on poll event, just signal ready, nothing else.
|
||
|
|
f.set_result(None)
|
||
|
|
return
|
||
|
|
elif kind == 'send_multipart':
|
||
|
|
send = self._shadow_sock.send_multipart
|
||
|
|
elif kind == 'send':
|
||
|
|
send = self._shadow_sock.send
|
||
|
|
else:
|
||
|
|
raise ValueError(f"Unhandled send event type: {kind!r}")
|
||
|
|
|
||
|
|
kwargs['flags'] |= _zmq.DONTWAIT
|
||
|
|
try:
|
||
|
|
result = send(msg, **kwargs)
|
||
|
|
except Exception as e:
|
||
|
|
f.set_exception(e)
|
||
|
|
else:
|
||
|
|
f.set_result(result)
|
||
|
|
|
||
|
|
# event masking from ZMQStream
|
||
|
|
def _handle_events(self, fd=0, events=0):
|
||
|
|
"""Dispatch IO events to _handle_recv, etc."""
|
||
|
|
if self._shadow_sock.closed:
|
||
|
|
return
|
||
|
|
|
||
|
|
zmq_events = self._shadow_sock.get(EVENTS)
|
||
|
|
if zmq_events & _zmq.POLLIN:
|
||
|
|
self._handle_recv()
|
||
|
|
if zmq_events & _zmq.POLLOUT:
|
||
|
|
self._handle_send()
|
||
|
|
self._schedule_remaining_events()
|
||
|
|
|
||
|
|
def _schedule_remaining_events(self, events=None):
|
||
|
|
"""Schedule a call to handle_events next loop iteration
|
||
|
|
|
||
|
|
If there are still events to handle.
|
||
|
|
"""
|
||
|
|
# edge-triggered handling
|
||
|
|
# allow passing events in, in case this is triggered by retrieving events,
|
||
|
|
# so we don't have to retrieve it twice.
|
||
|
|
if self._state == 0:
|
||
|
|
# not watching for anything, nothing to schedule
|
||
|
|
return
|
||
|
|
if events is None:
|
||
|
|
events = self._shadow_sock.get(EVENTS)
|
||
|
|
if events & self._state:
|
||
|
|
self._call_later(0, self._handle_events)
|
||
|
|
|
||
|
|
def _add_io_state(self, state):
|
||
|
|
"""Add io_state to poller."""
|
||
|
|
if self._state != state:
|
||
|
|
state = self._state = self._state | state
|
||
|
|
self._update_handler(self._state)
|
||
|
|
|
||
|
|
def _drop_io_state(self, state):
|
||
|
|
"""Stop poller from watching an io_state."""
|
||
|
|
if self._state & state:
|
||
|
|
self._state = self._state & (~state)
|
||
|
|
self._update_handler(self._state)
|
||
|
|
|
||
|
|
def _update_handler(self, state):
|
||
|
|
"""Update IOLoop handler with state.
|
||
|
|
|
||
|
|
zmq FD is always read-only.
|
||
|
|
"""
|
||
|
|
# ensure loop is registered and init_io has been called
|
||
|
|
# if there are any events to watch for
|
||
|
|
if state:
|
||
|
|
self._get_loop()
|
||
|
|
self._schedule_remaining_events()
|
||
|
|
|
||
|
|
def _init_io_state(self, loop=None):
|
||
|
|
"""initialize the ioloop event handler"""
|
||
|
|
if loop is None:
|
||
|
|
loop = self._get_loop()
|
||
|
|
loop.add_handler(self._shadow_sock, self._handle_events, self._READ)
|
||
|
|
self._call_later(0, self._handle_events)
|
||
|
|
|
||
|
|
def _clear_io_state(self):
|
||
|
|
"""unregister the ioloop event handler
|
||
|
|
|
||
|
|
called once during close
|
||
|
|
"""
|
||
|
|
fd = self._shadow_sock
|
||
|
|
if self._shadow_sock.closed:
|
||
|
|
fd = self._fd
|
||
|
|
if self._current_loop is not None:
|
||
|
|
self._current_loop.remove_handler(fd)
|