Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Derived from iostream.py from tornado 1.0, Copyright 2009 Facebook
2# Used under Apache License Version 2.0
3#
4# Modifications are Copyright (C) PyZMQ Developers
5# Distributed under the terms of the Modified BSD License.
6"""A utility class for event-based messaging on a zmq socket using tornado.
8.. seealso::
10 - :mod:`zmq.asyncio`
11 - :mod:`zmq.eventloop.future`
12"""
14from __future__ import annotations
16import asyncio
17import pickle
18import warnings
19from collections.abc import Awaitable, Sequence
20from queue import Queue
21from typing import TYPE_CHECKING, Any, Callable, Literal, overload
23from tornado.ioloop import IOLoop
24from tornado.log import gen_log
26import zmq
27import zmq._future
28from zmq import POLLIN, POLLOUT
29from zmq.utils import jsonapi
31if TYPE_CHECKING:
32 from typing_extensions import Concatenate, ParamSpec, Self, TypeAlias, TypeVar
34 from zmq.sugar.socket import _SocketContext
36 _RecvT = TypeVar("_RecvT", bound=bytes | zmq.Frame)
37 _P = ParamSpec("_P")
39 _Ignored: TypeAlias = object
41 _CallbackRecv: TypeAlias = Callable[[list[_RecvT]], _Ignored]
42 _CallbackRecvStream: TypeAlias = Callable[["ZMQStream", list[_RecvT]], _Ignored]
44 _CallbackSend: TypeAlias = Callable[
45 [list[Any], zmq.MessageTracker | None],
46 _Ignored,
47 ]
50class ZMQStream:
51 """A utility class to register callbacks when a zmq socket sends and receives
53 For use with tornado IOLoop.
55 There are three main methods
57 Methods:
59 * **on_recv(callback, copy=True):**
60 register a callback to be run every time the socket has something to receive
61 * **on_send(callback):**
62 register a callback to be run every time you call send
63 * **send_multipart(self, msg, flags=0, copy=False, callback=None):**
64 perform a send that will trigger the callback
65 if callback is passed, on_send is also called.
67 There are also send_multipart(), send_json(), send_pyobj()
69 Three other methods for deactivating the callbacks:
71 * **stop_on_recv():**
72 turn off the recv callback
73 * **stop_on_send():**
74 turn off the send callback
76 which simply call ``on_<evt>(None)``.
78 The entire socket interface, excluding direct recv methods, is also
79 provided, primarily through direct-linking the methods.
80 e.g.
82 >>> stream.bind is stream.socket.bind
83 True
86 .. versionadded:: 25
88 send/recv callbacks can be coroutines.
90 .. versionchanged:: 25
92 ZMQStreams only support base zmq.Socket classes (this has always been true, but not enforced).
93 If ZMQStreams are created with e.g. async Socket subclasses,
94 a RuntimeWarning will be shown,
95 and the socket cast back to the default zmq.Socket
96 before connecting events.
98 Previously, using async sockets (or any zmq.Socket subclass) would result in undefined behavior for the
99 arguments passed to callback functions.
100 Now, the callback functions reliably get the return value of the base `zmq.Socket` send/recv_multipart methods
101 (the list of message frames).
102 """
104 socket: zmq.Socket
105 io_loop: IOLoop
106 poller: zmq.Poller
107 _send_queue: Queue
108 _recv_callback: Callable | None
109 _send_callback: Callable | None
110 _close_callback: Callable | None
111 _state: int = 0
112 _flushed: bool = False
113 _recv_copy: bool = False
114 _fd: int
116 bind: Callable[[str], _SocketContext[zmq.Socket]]
117 bind_to_random_port: Callable[Concatenate[str, ...], int]
118 connect: Callable[[str], _SocketContext[zmq.Socket]]
119 setsockopt: Callable[[int, int | bytes], None]
120 getsockopt: Callable[[int], int | bytes]
121 setsockopt_string: Callable[Concatenate[int, str, ...], None]
122 setsockopt_unicode: Callable[Concatenate[int, str, ...], None]
123 getsockopt_string: Callable[Concatenate[int, ...], str]
124 getsockopt_unicode: Callable[Concatenate[int, ...], str]
126 def __init__(self, socket: zmq.Socket, io_loop: IOLoop | None = None) -> None:
127 if isinstance(socket, zmq._future._AsyncSocket):
128 warnings.warn(
129 f"""ZMQStream only supports the base zmq.Socket class.
131 Use zmq.Socket(shadow=other_socket)
132 or `ctx.socket(zmq.{socket._type_name}, socket_class=zmq.Socket)`
133 to create a base zmq.Socket object,
134 no matter what other kind of socket your Context creates.
135 """,
136 RuntimeWarning,
137 stacklevel=2,
138 )
139 # shadow back to base zmq.Socket,
140 # otherwise callbacks like `on_recv` will get the wrong types.
141 socket = zmq.Socket(shadow=socket)
142 self.socket = socket
144 # IOLoop.current() is deprecated if called outside the event loop
145 # that means
146 self.io_loop = io_loop or IOLoop.current()
147 self.poller = zmq.Poller()
148 self._fd = self.socket.FD
150 self._send_queue = Queue()
151 self._recv_callback = None
152 self._send_callback = None
153 self._close_callback = None
154 self._recv_copy = False
155 self._flushed = False
157 self._state = 0
158 self._init_io_state()
160 # shortcircuit some socket methods
161 self.bind = self.socket.bind
162 self.bind_to_random_port = self.socket.bind_to_random_port
163 self.connect = self.socket.connect
164 self.setsockopt = self.socket.setsockopt
165 self.getsockopt = self.socket.getsockopt
166 self.setsockopt_string = self.socket.setsockopt_string
167 self.getsockopt_string = self.socket.getsockopt_string
168 self.setsockopt_unicode = self.socket.setsockopt_unicode
169 self.getsockopt_unicode = self.socket.getsockopt_unicode
171 def stop_on_recv(self) -> None:
172 """Disable callback and automatic receiving."""
173 return self.on_recv(None)
175 def stop_on_send(self) -> None:
176 """Disable callback on sending."""
177 return self.on_send(None)
179 def stop_on_err(self) -> None:
180 """DEPRECATED, does nothing"""
181 gen_log.warning("on_err does nothing, and will be removed")
183 def on_err(self, callback: Callable) -> None:
184 """DEPRECATED, does nothing"""
185 gen_log.warning("on_err does nothing, and will be removed")
187 @overload
188 def on_recv(
189 self,
190 callback: _CallbackRecv[bytes] | None,
191 copy: Literal[True] = True,
192 ) -> None: ...
193 @overload
194 def on_recv(
195 self,
196 callback: _CallbackRecv[zmq.Frame],
197 copy: Literal[False],
198 ) -> None: ...
199 @overload
200 def on_recv(
201 self,
202 callback: _CallbackRecv[_RecvT] | None,
203 copy: bool = True,
204 ) -> None: ...
205 def on_recv(
206 self,
207 callback: _CallbackRecv[_RecvT] | None,
208 copy: bool = True,
209 ) -> None:
210 """Register a callback for when a message is ready to recv.
212 There can be only one callback registered at a time, so each
213 call to `on_recv` replaces previously registered callbacks.
215 on_recv(None) disables recv event polling.
217 Use on_recv_stream(callback) instead, to register a callback that will receive
218 both this ZMQStream and the message, instead of just the message.
220 Parameters
221 ----------
223 callback : callable
224 callback must take exactly one argument, which will be a
225 list, as returned by socket.recv_multipart()
226 if callback is None, recv callbacks are disabled.
227 copy : bool
228 copy is passed directly to recv, so if copy is False,
229 callback will receive Message objects. If copy is True,
230 then callback will receive bytes/str objects.
232 Returns : None
233 """
235 self._check_closed()
236 assert callback is None or callable(callback)
237 self._recv_callback = callback
238 self._recv_copy = copy
239 if callback is None:
240 self._drop_io_state(zmq.POLLIN)
241 else:
242 self._add_io_state(zmq.POLLIN)
244 @overload
245 def on_recv_stream(
246 self,
247 callback: _CallbackRecvStream[bytes] | None,
248 copy: Literal[True] = True,
249 ) -> None: ...
250 @overload
251 def on_recv_stream(
252 self,
253 callback: _CallbackRecvStream[zmq.Frame],
254 copy: Literal[False],
255 ) -> None: ...
256 @overload
257 def on_recv_stream(
258 self,
259 callback: _CallbackRecvStream[_RecvT] | None,
260 copy: bool = True,
261 ) -> None: ...
262 def on_recv_stream(
263 self,
264 callback: _CallbackRecvStream[_RecvT] | None,
265 copy: bool = True,
266 ) -> None:
267 """Same as on_recv, but callback will get this stream as first argument
269 callback must take exactly two arguments, as it will be called as::
271 callback(stream, msg)
273 Useful when a single callback should be used with multiple streams.
274 """
275 if callback is None:
276 self.stop_on_recv()
277 else:
279 def stream_callback(msg):
280 return callback(self, msg)
282 self.on_recv(stream_callback, copy=copy)
284 def on_send(self, callback: _CallbackSend | None) -> None:
285 """Register a callback to be called on each send
287 There will be two arguments::
289 callback(msg, status)
291 * `msg` will be the list of sendable objects that was just sent
292 * `status` will be the return result of socket.send_multipart(msg) -
293 MessageTracker or None.
295 Non-copying sends return a MessageTracker object whose
296 `done` attribute will be True when the send is complete.
297 This allows users to track when an object is safe to write to
298 again.
300 The second argument will always be None if copy=True
301 on the send.
303 Use on_send_stream(callback) to register a callback that will be passed
304 this ZMQStream as the first argument, in addition to the other two.
306 on_send(None) disables recv event polling.
308 Parameters
309 ----------
311 callback : callable
312 callback must take exactly two arguments, which will be
313 the message being sent (always a list),
314 and the return result of socket.send_multipart(msg) -
315 MessageTracker or None.
317 if callback is None, send callbacks are disabled.
318 """
320 self._check_closed()
321 assert callback is None or callable(callback)
322 self._send_callback = callback
324 def on_send_stream(
325 self,
326 callback: Callable[[Self, list[Any], zmq.MessageTracker | None], Any] | None,
327 ) -> None:
328 """Same as on_send, but callback will get this stream as first argument
330 Callback will be passed three arguments::
332 callback(stream, msg, status)
334 Useful when a single callback should be used with multiple streams.
335 """
336 if callback is None:
337 self.stop_on_send()
338 else:
339 self.on_send(lambda msg, status: callback(self, msg, status))
341 def send(
342 self,
343 msg: Any,
344 flags: int = 0,
345 copy: bool = True,
346 track: bool = False,
347 callback: _CallbackSend | None = None,
348 **kwargs: Any,
349 ) -> None:
350 """Send a message, optionally also register a new callback for sends.
351 See zmq.socket.send for details.
352 """
353 return self.send_multipart(
354 [msg], flags=flags, copy=copy, track=track, callback=callback, **kwargs
355 )
357 def send_multipart(
358 self,
359 msg: Sequence[Any],
360 flags: int = 0,
361 copy: bool = True,
362 track: bool = False,
363 callback: _CallbackSend | None = None,
364 **kwargs: Any,
365 ) -> None:
366 """Send a multipart message, optionally also register a new callback for sends.
367 See zmq.socket.send_multipart for details.
368 """
369 kwargs.update(dict(flags=flags, copy=copy, track=track))
370 self._send_queue.put((msg, kwargs))
371 callback = callback or self._send_callback
372 if callback is not None:
373 self.on_send(callback)
374 else:
375 # noop callback
376 self.on_send(lambda *args: None)
377 self._add_io_state(zmq.POLLOUT)
379 def send_string(
380 self,
381 u: str,
382 flags: int = 0,
383 encoding: str = 'utf-8',
384 callback: _CallbackSend | None = None,
385 **kwargs: Any,
386 ) -> None:
387 """Send a unicode message with an encoding.
388 See zmq.socket.send_unicode for details.
389 """
390 if not isinstance(u, str):
391 raise TypeError("unicode/str objects only")
392 return self.send(u.encode(encoding), flags=flags, callback=callback, **kwargs)
394 send_unicode = send_string
396 def send_json(
397 self,
398 obj: Any,
399 flags: int = 0,
400 callback: _CallbackSend | None = None,
401 **kwargs: Any,
402 ) -> None:
403 """Send json-serialized version of an object.
404 See zmq.socket.send_json for details.
405 """
406 msg = jsonapi.dumps(obj)
407 return self.send(msg, flags=flags, callback=callback, **kwargs)
409 def send_pyobj(
410 self,
411 obj: Any,
412 flags: int = 0,
413 protocol: int = -1,
414 callback: _CallbackSend | None = None,
415 **kwargs: Any,
416 ) -> None:
417 """Send a Python object as a message using pickle to serialize.
419 See zmq.socket.send_json for details.
420 """
421 msg = pickle.dumps(obj, protocol)
422 return self.send(msg, flags, callback=callback, **kwargs)
424 def _finish_flush(self) -> None:
425 """callback for unsetting _flushed flag."""
426 self._flushed = False
428 def flush(
429 self,
430 flag: int = zmq.POLLIN | zmq.POLLOUT,
431 limit: int | None = None,
432 ) -> int:
433 """Flush pending messages.
435 This method safely handles all pending incoming and/or outgoing messages,
436 bypassing the inner loop, passing them to the registered callbacks.
438 A limit can be specified, to prevent blocking under high load.
440 flush will return the first time ANY of these conditions are met:
441 * No more events matching the flag are pending.
442 * the total number of events handled reaches the limit.
444 Note that if ``flag|POLLIN != 0``, recv events will be flushed even if no callback
445 is registered, unlike normal IOLoop operation. This allows flush to be
446 used to remove *and ignore* incoming messages.
448 Parameters
449 ----------
450 flag : int
451 default=POLLIN|POLLOUT
452 0MQ poll flags.
453 If flag|POLLIN, recv events will be flushed.
454 If flag|POLLOUT, send events will be flushed.
455 Both flags can be set at once, which is the default.
456 limit : None or int, optional
457 The maximum number of messages to send or receive.
458 Both send and recv count against this limit.
460 Returns
461 -------
462 int :
463 count of events handled (both send and recv)
464 """
465 self._check_closed()
466 # unset self._flushed, so callbacks will execute, in case flush has
467 # already been called this iteration
468 already_flushed = self._flushed
469 self._flushed = False
470 # initialize counters
471 count = 0
473 def update_flag():
474 """Update the poll flag, to prevent registering POLLOUT events
475 if we don't have pending sends."""
476 return flag & zmq.POLLIN | (self.sending() and flag & zmq.POLLOUT)
478 flag = update_flag()
479 if not flag:
480 # nothing to do
481 return 0
482 self.poller.register(self.socket, flag)
483 events = self.poller.poll(0)
484 while events and (not limit or count < limit):
485 s, event = events[0]
486 if event & POLLIN: # receiving
487 self._handle_recv()
488 count += 1
489 if self.socket is None:
490 # break if socket was closed during callback
491 break
492 if event & POLLOUT and self.sending():
493 self._handle_send()
494 count += 1
495 if self.socket is None:
496 # break if socket was closed during callback
497 break
499 flag = update_flag()
500 if flag:
501 self.poller.register(self.socket, flag)
502 events = self.poller.poll(0)
503 else:
504 events = []
505 if count: # only bypass loop if we actually flushed something
506 # skip send/recv callbacks this iteration
507 self._flushed = True
508 # reregister them at the end of the loop
509 if not already_flushed: # don't need to do it again
510 self.io_loop.add_callback(self._finish_flush)
511 elif already_flushed:
512 self._flushed = True
514 # update ioloop poll state, which may have changed
515 self._rebuild_io_state()
516 return count
518 def set_close_callback(self, callback: Callable[[], _Ignored] | None) -> None:
519 """Call the given callback when the stream is closed."""
520 self._close_callback = callback
522 def close(self, linger: int | None = None) -> None:
523 """Close this stream."""
524 if self.socket is not None:
525 if self.socket.closed:
526 # fallback on raw fd for closed sockets
527 # hopefully this happened promptly after close,
528 # otherwise somebody else may have the FD
529 warnings.warn(
530 f"Unregistering FD {self._fd} after closing socket. "
531 "This could result in unregistering handlers for the wrong socket. "
532 "Please use stream.close() instead of closing the socket directly.",
533 stacklevel=2,
534 )
535 self.io_loop.remove_handler(self._fd)
536 else:
537 self.io_loop.remove_handler(self.socket)
538 self.socket.close(linger)
539 self.socket = None # type: ignore
540 if self._close_callback:
541 self._run_callback(self._close_callback)
543 def receiving(self) -> bool:
544 """Returns True if we are currently receiving from the stream."""
545 return self._recv_callback is not None
547 def sending(self) -> bool:
548 """Returns True if we are currently sending to the stream."""
549 return not self._send_queue.empty()
551 def closed(self) -> bool:
552 if self.socket is None:
553 return True
554 if self.socket.closed:
555 # underlying socket has been closed, but not by us!
556 # trigger our cleanup
557 self.close()
558 return True
559 return False
561 def _run_callback(
562 self,
563 callback: Callable[_P, Any],
564 *args: _P.args,
565 **kwargs: _P.kwargs,
566 ) -> None:
567 """Wrap running callbacks in try/except to allow us to
568 close our socket."""
569 try:
570 f = callback(*args, **kwargs)
571 if isinstance(f, Awaitable):
572 f = asyncio.ensure_future(f)
573 else:
574 f = None
575 except Exception:
576 gen_log.error("Uncaught exception in ZMQStream callback", exc_info=True)
577 # Re-raise the exception so that IOLoop.handle_callback_exception
578 # can see it and log the error
579 raise
581 if f is not None:
582 # handle async callbacks
583 def _log_error(f):
584 try:
585 f.result()
586 except Exception:
587 gen_log.error(
588 "Uncaught exception in ZMQStream callback", exc_info=True
589 )
591 f.add_done_callback(_log_error)
593 def _handle_events(self, fd, events):
594 """This method is the actual handler for IOLoop, that gets called whenever
595 an event on my socket is posted. It dispatches to _handle_recv, etc."""
596 if not self.socket:
597 gen_log.warning("Got events for closed stream %s", self)
598 return
599 try:
600 zmq_events = self.socket.EVENTS
601 except zmq.ContextTerminated:
602 gen_log.warning("Got events for stream %s after terminating context", self)
603 # trigger close check, this will unregister callbacks
604 self.closed()
605 return
606 except zmq.ZMQError as e:
607 # run close check
608 # shadow sockets may have been closed elsewhere,
609 # which should show up as ENOTSOCK here
610 if self.closed():
611 gen_log.warning(
612 "Got events for stream %s attached to closed socket: %s", self, e
613 )
614 else:
615 gen_log.error("Error getting events for %s: %s", self, e)
616 return
617 try:
618 # dispatch events:
619 if zmq_events & zmq.POLLIN and self.receiving():
620 self._handle_recv()
621 if not self.socket:
622 return
623 if zmq_events & zmq.POLLOUT and self.sending():
624 self._handle_send()
625 if not self.socket:
626 return
628 # rebuild the poll state
629 self._rebuild_io_state()
630 except Exception:
631 gen_log.error("Uncaught exception in zmqstream callback", exc_info=True)
632 raise
634 def _handle_recv(self):
635 """Handle a recv event."""
636 if self._flushed:
637 return
638 try:
639 msg = self.socket.recv_multipart(zmq.NOBLOCK, copy=self._recv_copy)
640 except zmq.ZMQError as e:
641 if e.errno == zmq.EAGAIN:
642 # state changed since poll event
643 pass
644 else:
645 raise
646 else:
647 if self._recv_callback:
648 callback = self._recv_callback
649 self._run_callback(callback, msg)
651 def _handle_send(self):
652 """Handle a send event."""
653 if self._flushed:
654 return
655 if not self.sending():
656 gen_log.error("Shouldn't have handled a send event")
657 return
659 msg, kwargs = self._send_queue.get()
660 try:
661 status = self.socket.send_multipart(msg, **kwargs)
662 except zmq.ZMQError as e:
663 gen_log.error("SEND Error: %s", e)
664 status = e
665 if self._send_callback:
666 callback = self._send_callback
667 self._run_callback(callback, msg, status)
669 def _check_closed(self):
670 if not self.socket:
671 raise OSError("Stream is closed")
673 def _rebuild_io_state(self):
674 """rebuild io state based on self.sending() and receiving()"""
675 if self.socket is None:
676 return
677 state = 0
678 if self.receiving():
679 state |= zmq.POLLIN
680 if self.sending():
681 state |= zmq.POLLOUT
683 self._state = state
684 self._update_handler(state)
686 def _add_io_state(self, state):
687 """Add io_state to poller."""
688 self._state = self._state | state
689 self._update_handler(self._state)
691 def _drop_io_state(self, state):
692 """Stop poller from watching an io_state."""
693 self._state = self._state & (~state)
694 self._update_handler(self._state)
696 def _update_handler(self, state):
697 """Update IOLoop handler with state."""
698 if self.socket is None:
699 return
701 if state & self.socket.events:
702 # events still exist that haven't been processed
703 # explicitly schedule handling to avoid missing events due to edge-triggered FDs
704 self.io_loop.add_callback(lambda: self._handle_events(self.socket, 0))
706 def _init_io_state(self):
707 """initialize the ioloop event handler"""
708 self.io_loop.add_handler(self.socket, self._handle_events, self.io_loop.READ)