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"""Base classes to manage a Client's interaction with a running kernel"""
3# Copyright (c) Jupyter Development Team.
4# Distributed under the terms of the Modified BSD License.
5import asyncio
6import atexit
7import time
8import typing as t
9from queue import Empty
10from threading import Event, Thread
12import zmq.asyncio
13from jupyter_core.utils import ensure_async
15from ._version import protocol_version_info
16from .channelsabc import HBChannelABC
17from .session import Session
19# import ZMQError in top-level namespace, to avoid ugly attribute-error messages
20# during garbage collection of threads at exit
22# -----------------------------------------------------------------------------
23# Constants and exceptions
24# -----------------------------------------------------------------------------
26major_protocol_version = protocol_version_info[0]
29class InvalidPortNumber(Exception): # noqa
30 """An exception raised for an invalid port number."""
32 pass
35class HBChannel(Thread):
36 """The heartbeat channel which monitors the kernel heartbeat.
38 Note that the heartbeat channel is paused by default. As long as you start
39 this channel, the kernel manager will ensure that it is paused and un-paused
40 as appropriate.
41 """
43 session = None
44 socket = None
45 address = None
46 _exiting = False
48 time_to_dead: float = 1.0
49 _running = None
50 _pause = None
51 _beating = None
53 def __init__(
54 self,
55 context: zmq.Context | None = None,
56 session: Session | None = None,
57 address: t.Union[t.Tuple[str, int], str] = "",
58 *,
59 curve_serverkey: bytes | None = None,
60 ) -> None:
61 """Create the heartbeat monitor thread.
63 Parameters
64 ----------
65 context : :class:`zmq.Context`
66 The ZMQ context to use.
67 session : :class:`session.Session`
68 The session to use.
69 address : zmq url
70 Standard (ip, port) tuple that the kernel is listening on.
71 curve_serverkey : bytes, optional
72 CurveZMQ server public key (Z85). When provided, the
73 heartbeat REQ socket is configured as a CurveZMQ client so it
74 can communicate with a CurveZMQ-enabled kernel.
75 """
76 super().__init__()
77 self.daemon = True
79 self.context = context
80 self.session = session
81 self.curve_serverkey = curve_serverkey
82 if isinstance(address, tuple):
83 if address[1] == 0:
84 message = "The port number for a channel cannot be 0."
85 raise InvalidPortNumber(message)
86 address_str = "tcp://%s:%i" % address
87 else:
88 address_str = address
89 self.address = address_str
91 # running is False until `.start()` is called
92 self._running = False
93 self._exit = Event()
94 # don't start paused
95 self._pause = False
96 self.poller = zmq.Poller()
98 @staticmethod
99 @atexit.register
100 def _notice_exit() -> None:
101 # Class definitions can be torn down during interpreter shutdown.
102 # We only need to set _exiting flag if this hasn't happened.
103 if HBChannel is not None:
104 HBChannel._exiting = True
106 def _create_socket(self) -> None:
107 if self.socket is not None:
108 # close previous socket, before opening a new one
109 self.poller.unregister(self.socket) # type:ignore[unreachable]
110 self.socket.close()
111 assert self.context is not None
112 self.socket = self.context.socket(zmq.REQ)
113 self.socket.linger = 1000
114 if self.curve_serverkey is not None:
115 # Generate a fresh ephemeral keypair for each socket; only the
116 # server public key (curve_serverkey) is needed for authentication.
117 client_pub, client_sec = zmq.curve_keypair()
118 self.socket.curve_secretkey = client_sec
119 self.socket.curve_publickey = client_pub
120 self.socket.curve_serverkey = self.curve_serverkey
121 assert self.address is not None
122 self.socket.connect(self.address)
124 self.poller.register(self.socket, zmq.POLLIN)
126 async def _async_run(self) -> None:
127 """The thread's main activity. Call start() instead."""
128 self._create_socket()
129 self._running = True
130 self._beating = True
131 assert self.socket is not None
133 while self._running:
134 if self._pause:
135 # just sleep, and skip the rest of the loop
136 self._exit.wait(self.time_to_dead)
137 continue
139 since_last_heartbeat = 0.0
140 # no need to catch EFSM here, because the previous event was
141 # either a recv or connect, which cannot be followed by EFSM)
142 await ensure_async(self.socket.send(b"ping"))
143 request_time = time.time()
144 # Wait until timeout
145 self._exit.wait(self.time_to_dead)
146 # poll(0) means return immediately (see http://api.zeromq.org/2-1:zmq-poll)
147 self._beating = bool(self.poller.poll(0))
148 if self._beating:
149 # the poll above guarantees we have something to recv
150 await ensure_async(self.socket.recv())
151 continue
152 elif self._running:
153 # nothing was received within the time limit, signal heart failure
154 since_last_heartbeat = time.time() - request_time
155 self.call_handlers(since_last_heartbeat)
156 # and close/reopen the socket, because the REQ/REP cycle has been broken
157 self._create_socket()
158 continue
160 def run(self) -> None:
161 """Run the heartbeat thread."""
162 loop = asyncio.new_event_loop()
163 asyncio.set_event_loop(loop)
164 try:
165 loop.run_until_complete(self._async_run())
166 finally:
167 loop.close()
169 def pause(self) -> None:
170 """Pause the heartbeat."""
171 self._pause = True
173 def unpause(self) -> None:
174 """Unpause the heartbeat."""
175 self._pause = False
177 def is_beating(self) -> bool:
178 """Is the heartbeat running and responsive (and not paused)."""
179 if self.is_alive() and not self._pause and self._beating: # noqa
180 return True
181 else:
182 return False
184 def stop(self) -> None:
185 """Stop the channel's event loop and join its thread."""
186 self._running = False
187 self._exit.set()
188 self.join()
189 self.close()
191 def close(self) -> None:
192 """Close the heartbeat thread."""
193 if self.socket is not None:
194 try:
195 self.socket.close(linger=0)
196 except Exception:
197 pass
198 self.socket = None
200 def call_handlers(self, since_last_heartbeat: float) -> None:
201 """This method is called in the ioloop thread when a message arrives.
203 Subclasses should override this method to handle incoming messages.
204 It is important to remember that this method is called in the thread
205 so that some logic must be done to ensure that the application level
206 handlers are called in the application thread.
207 """
208 pass
211HBChannelABC.register(HBChannel)
214class ZMQSocketChannel:
215 """A ZMQ socket wrapper"""
217 def __init__(self, socket: zmq.Socket, session: Session, loop: t.Any = None) -> None:
218 """Create a channel.
220 Parameters
221 ----------
222 socket : :class:`zmq.Socket`
223 The ZMQ socket to use.
224 session : :class:`session.Session`
225 The session to use.
226 loop
227 Unused here, for other implementations
228 """
229 super().__init__()
231 self.socket: zmq.Socket | None = socket
232 self.session = session
234 def _recv(self, **kwargs: t.Any) -> t.Dict[str, t.Any]:
235 assert self.socket is not None
236 msg = self.socket.recv_multipart(**kwargs)
237 _ident, smsg = self.session.feed_identities(msg)
238 return self.session.deserialize(smsg)
240 def get_msg(self, timeout: float | None = None) -> t.Dict[str, t.Any]:
241 """Gets a message if there is one that is ready."""
242 assert self.socket is not None
243 timeout_ms = None if timeout is None else int(timeout * 1000) # seconds to ms
244 ready = self.socket.poll(timeout_ms)
245 if ready:
246 res = self._recv()
247 return res
248 else:
249 raise Empty
251 def get_msgs(self) -> t.List[t.Dict[str, t.Any]]:
252 """Get all messages that are currently ready."""
253 msgs = []
254 while True:
255 try:
256 msgs.append(self.get_msg())
257 except Empty:
258 break
259 return msgs
261 def msg_ready(self) -> bool:
262 """Is there a message that has been received?"""
263 assert self.socket is not None
264 return bool(self.socket.poll(timeout=0))
266 def close(self) -> None:
267 """Close the socket channel."""
268 if self.socket is not None:
269 try:
270 self.socket.close(linger=0)
271 except Exception:
272 pass
273 self.socket = None
275 stop = close
277 def is_alive(self) -> bool:
278 """Test whether the channel is alive."""
279 return self.socket is not None
281 def send(self, msg: t.Dict[str, t.Any]) -> None:
282 """Pass a message to the ZMQ socket to send"""
283 assert self.socket is not None
284 self.session.send(self.socket, msg)
286 def start(self) -> None:
287 """Start the socket channel."""
288 pass
291class AsyncZMQSocketChannel(ZMQSocketChannel):
292 """A ZMQ socket in an async API"""
294 socket: zmq.asyncio.Socket
296 def __init__(self, socket: zmq.asyncio.Socket, session: Session, loop: t.Any = None) -> None:
297 """Create a channel.
299 Parameters
300 ----------
301 socket : :class:`zmq.asyncio.Socket`
302 The ZMQ socket to use.
303 session : :class:`session.Session`
304 The session to use.
305 loop
306 Unused here, for other implementations
307 """
308 if not isinstance(socket, zmq.asyncio.Socket):
309 msg = "Socket must be asyncio" # type:ignore[unreachable]
310 raise ValueError(msg)
311 super().__init__(socket, session)
313 async def _recv(self, **kwargs: t.Any) -> t.Dict[str, t.Any]: # type:ignore[override]
314 assert self.socket is not None
315 msg = await self.socket.recv_multipart(**kwargs)
316 _, smsg = self.session.feed_identities(msg)
317 return self.session.deserialize(smsg)
319 async def get_msg( # type:ignore[override]
320 self, timeout: float | None = None
321 ) -> t.Dict[str, t.Any]:
322 """Gets a message if there is one that is ready."""
323 assert self.socket is not None
324 timeout_ms = None if timeout is None else int(timeout * 1000) # seconds to ms
325 ready = await self.socket.poll(timeout_ms)
326 if ready:
327 res = await self._recv()
328 return res
329 else:
330 raise Empty
332 async def get_msgs(self) -> t.List[t.Dict[str, t.Any]]: # type:ignore[override]
333 """Get all messages that are currently ready."""
334 msgs = []
335 while True:
336 try:
337 msgs.append(await self.get_msg())
338 except Empty:
339 break
340 return msgs
342 async def msg_ready(self) -> bool: # type:ignore[override]
343 """Is there a message that has been received?"""
344 assert self.socket is not None
345 return bool(await self.socket.poll(timeout=0))