Merge pull request #279 from rolweber/port-already-in-use

Port already in use

Author: minrk

jupyter_client/connect.py

@@ -333,6 +333,9 @@ def _ip_changed(self, name, old, new):
     control_port = Integer(0, config=True,
             help="set the control (ROUTER) port [default: random]")
 
+    # names of the ports with random assignment
+    _random_port_names = None
+
     @property
     def ports(self):
         return [ getattr(self, name) for name in port_names ]
@@ -417,6 +420,37 @@ def cleanup_ipc_files(self):
             except (IOError, OSError):
                 pass
 
+    def _record_random_port_names(self):
+        """Records which of the ports are randomly assigned.
+
+        Records on first invocation, if the transport is tcp.
+        Does nothing on later invocations."""
+
+        if self.transport != 'tcp':
+            return
+        if self._random_port_names is not None:
+            return
+
+        self._random_port_names = []
+        for name in port_names:
+            if getattr(self, name) <= 0:
+                self._random_port_names.append(name)
+
+    def cleanup_random_ports(self):
+        """Forgets randomly assigned port numbers and cleans up the connection file.
+
+        Does nothing if no port numbers have been randomly assigned.
+        In particular, does nothing unless the transport is tcp.
+        """
+
+        if not self._random_port_names:
+            return
+
+        for name in self._random_port_names:
+            setattr(self, name, 0)
+
+        self.cleanup_connection_file()
+
     def write_connection_file(self):
         """Write connection info to JSON dict in self.connection_file."""
         if self._connection_file_written and os.path.exists(self.connection_file):
@@ -431,6 +465,7 @@ def write_connection_file(self):
             kernel_name=self.kernel_name
         )
         # write_connection_file also sets default ports:
+        self._record_random_port_names()
         for name in port_names:
             setattr(self, name, cfg[name])
 
@@ -467,6 +502,7 @@ def load_connection_info(self, info):
         self.transport = info.get('transport', self.transport)
         self.ip = info.get('ip', self._ip_default())
 
+        self._record_random_port_names()
         for name in port_names:
             if getattr(self, name) == 0 and name in info:
                 # not overridden by config or cl_args

jupyter_client/manager.py

@@ -326,12 +326,9 @@ def shutdown_kernel(self, now=False, restart=False):
 
         self.cleanup(connection_file=not restart)
 
-    def restart_kernel(self, now=False, **kw):
+    def restart_kernel(self, now=False, newports=False, **kw):
         """Restarts a kernel with the arguments that were used to launch it.
 
-        If the old kernel was launched with random ports, the same ports will be
-        used for the new kernel. The same connection file is used again.
-
         Parameters
         ----------
         now : bool, optional
@@ -342,6 +339,14 @@ def restart_kernel(self, now=False, **kw):
             In all cases the kernel is restarted, the only difference is whether
             it is given a chance to perform a clean shutdown or not.
 
+        newports : bool, optional
+            If the old kernel was launched with random ports, this flag decides
+            whether the same ports and connection file will be used again.
+            If False, the same ports and connection file are used. This is
+            the default. If True, new random port numbers are chosen and a
+            new connection file is written. It is still possible that the newly
+            chosen random port numbers happen to be the same as the old ones.
+
         `**kw` : optional
             Any options specified here will overwrite those used to launch the
             kernel.
@@ -353,6 +358,9 @@ def restart_kernel(self, now=False, **kw):
             # Stop currently running kernel.
             self.shutdown_kernel(now=now, restart=True)
 
+            if newports:
+                self.cleanup_random_ports()
+
             # Start new kernel.
             self._launch_args.update(kw)
             self.start_kernel(**self._launch_args)

jupyter_client/restarter.py

@@ -34,8 +34,13 @@ class KernelRestarter(LoggingConfigurable):
     restart_limit = Integer(5, config=True,
         help="""The number of consecutive autorestarts before the kernel is presumed dead."""
     )
+
+    random_ports_until_alive = Bool(True, config=True,
+        help="""Whether to choose new random ports when restarting before the kernel is alive."""
+    )
     _restarting = Bool(False)
     _restart_count = Integer(0)
+    _initial_startup = Bool(True)
 
     callbacks = Dict()
     def _callbacks_default(self):
@@ -98,14 +103,18 @@ def poll(self):
                 self._restart_count = 0
                 self.stop()
             else:
-                self.log.info('KernelRestarter: restarting kernel (%i/%i)',
+                newports = self.random_ports_until_alive and self._initial_startup
+                self.log.info('KernelRestarter: restarting kernel (%i/%i), %s random ports',
                     self._restart_count,
-                    self.restart_limit
+                    self.restart_limit,
+                    'new' if newports else 'keep'
                 )
                 self._fire_callbacks('restart')
-                self.kernel_manager.restart_kernel(now=True)
+                self.kernel_manager.restart_kernel(now=True, newports=newports)
                 self._restarting = True
         else:
+            if self._initial_startup:
+                self._initial_startup = False
             if self._restarting:
                 self.log.debug("KernelRestarter: restart apparently succeeded")
             self._restarting = False

jupyter_client/tests/test_connect.py

@@ -21,6 +21,10 @@ def initialize(self, argv=[]):
         JupyterApp.initialize(self, argv=argv)
         self.init_connection_file()
 
+class DummyConfigurable(connect.ConnectionFileMixin):
+    def initialize(self):
+        pass
+
 sample_info = dict(ip='1.2.3.4', transport='ipc',
         shell_port=1, hb_port=2, iopub_port=3, stdin_port=4, control_port=5,
         key=b'abc123', signature_scheme='hmac-md5', kernel_name='python'
@@ -31,6 +35,7 @@ def initialize(self, argv=[]):
         key=b'abc123', signature_scheme='hmac-md5', kernel_name='test'
     )
 
+
 def test_write_connection_file():
     with TemporaryDirectory() as d:
         cf = os.path.join(d, 'kernel.json')
@@ -171,4 +176,26 @@ def test_find_connection_file_abspath():
         with open(cf, 'w') as f:
             f.write('{}')
         assert connect.find_connection_file(abs_cf, path=jupyter_runtime_dir()) == abs_cf
+        os.remove(abs_cf)
+
+
+def test_mixin_record_random_ports():
+    with TemporaryDirectory() as d:
+        dc = DummyConfigurable(data_dir=d, kernel_name='via-tcp', transport='tcp')
+        dc.write_connection_file()
+
+        assert dc._connection_file_written
+        assert os.path.exists(dc.connection_file)
+        assert dc._random_port_names == connect.port_names
 
+
+def test_mixin_cleanup_random_ports():
+    with TemporaryDirectory() as d:
+        dc = DummyConfigurable(data_dir=d, kernel_name='via-tcp', transport='tcp')
+        dc.write_connection_file()
+        filename = dc.connection_file
+        dc.cleanup_random_ports()
+
+        assert not os.path.exists(filename)
+        for name in dc._random_port_names:
+            assert getattr(dc, name) == 0