1  """The client and server for a basic ping-pong style heartbeat.
   2  """
   3  
   4  # -----------------------------------------------------------------------------
   5  #  Copyright (C) 2008-2011  The IPython Development Team
   6  #
   7  #  Distributed under the terms of the BSD License.  The full license is in
   8  #  the file LICENSE, distributed as part of this software.
   9  # -----------------------------------------------------------------------------
  10  
  11  # -----------------------------------------------------------------------------
  12  # Imports
  13  # -----------------------------------------------------------------------------
  14  
  15  import errno
  16  import socket
  17  from pathlib import Path
  18  from threading import Thread
  19  
  20  import zmq
  21  from jupyter_client.localinterfaces import localhost
  22  
  23  # -----------------------------------------------------------------------------
  24  # Code
  25  # -----------------------------------------------------------------------------
  26  
  27  
  28  class Heartbeat(Thread):
  29      """A simple ping-pong style heartbeat that runs in a thread."""
  30  
  31      def __init__(self, context, addr=None):
  32          """Initialize the heartbeat thread."""
  33          if addr is None:
  34              addr = ("tcp", localhost(), 0)
  35          Thread.__init__(self, name="Heartbeat")
  36          self.context = context
  37          self.transport, self.ip, self.port = addr
  38          self.original_port = self.port
  39          if self.original_port == 0:
  40              self.pick_port()
  41          self.addr = (self.ip, self.port)
  42          self.daemon = True
  43          self.pydev_do_not_trace = True
  44          self.is_pydev_daemon_thread = True
  45          self.name = "Heartbeat"
  46  
  47      def pick_port(self):
  48          """Pick a port for the heartbeat."""
  49          if self.transport == "tcp":
  50              s = socket.socket()
  51              # '*' means all interfaces to 0MQ, which is '' to socket.socket
  52              s.bind(("" if self.ip == "*" else self.ip, 0))
  53              self.port = s.getsockname()[1]
  54              s.close()
  55          elif self.transport == "ipc":
  56              self.port = 1
  57              while Path(f"{self.ip}-{self.port}").exists():
  58                  self.port = self.port + 1
  59          else:
  60              raise ValueError("Unrecognized zmq transport: %s" % self.transport)
  61          return self.port
  62  
  63      def _try_bind_socket(self):
  64          c = ":" if self.transport == "tcp" else "-"
  65          return self.socket.bind(f"{self.transport}://{self.ip}" + c + str(self.port))
  66  
  67      def _bind_socket(self):
  68          try:
  69              win_in_use = errno.WSAEADDRINUSE  # type:ignore[attr-defined]
  70          except AttributeError:
  71              win_in_use = None
  72  
  73          # Try up to 100 times to bind a port when in conflict to avoid
  74          # infinite attempts in bad setups
  75          max_attempts = 1 if self.original_port else 100
  76          for attempt in range(max_attempts):
  77              try:
  78                  self._try_bind_socket()
  79              except zmq.ZMQError as ze:
  80                  if attempt == max_attempts - 1:
  81                      raise
  82                  # Raise if we have any error not related to socket binding
  83                  if ze.errno != errno.EADDRINUSE and ze.errno != win_in_use:
  84                      raise
  85                  # Raise if we have any error not related to socket binding
  86                  if self.original_port == 0:
  87                      self.pick_port()
  88                  else:
  89                      raise
  90              else:
  91                  return
  92  
  93      def run(self):
  94          """Run the heartbeat thread."""
  95          self.name = "Heartbeat"
  96          self.socket = self.context.socket(zmq.ROUTER)
  97          self.socket.linger = 1000
  98          try:
  99              self._bind_socket()
 100          except Exception:
 101              self.socket.close()
 102              raise
 103  
 104          while True:
 105              try:
 106                  zmq.device(zmq.QUEUE, self.socket, self.socket)
 107              except zmq.ZMQError as e:
 108                  if e.errno == errno.EINTR:
 109                      # signal interrupt, resume heartbeat
 110                      continue
 111                  if e.errno == zmq.ETERM:
 112                      # context terminated, close socket and exit
 113                      try:
 114                          self.socket.close()
 115                      except zmq.ZMQError:
 116                          # suppress further errors during cleanup
 117                          # this shouldn't happen, though
 118                          pass
 119                      break
 120                  if e.errno == zmq.ENOTSOCK:
 121                      # socket closed elsewhere, exit
 122                      break
 123                  raise
 124              else:
 125                  break
 126