Fixed subtle bug in kernel restarting.

* Old routers were not being shutdown and removed.
* We were incorrectly associating the new kernel with the notebook
  (we were using the *old* kernel_id for this).
* General clean ups in the kernel manager.

IPython/frontend/html/notebook/handlers.py

@@ -1,160 +1,158 b''
 """Tornado handlers for the notebook."""
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import json
 import logging
 import urllib
 
 from tornado import web
 from tornado import websocket
 
 
 #-----------------------------------------------------------------------------
 # Top-level handlers
 #-----------------------------------------------------------------------------
 
 
 class NBBrowserHandler(web.RequestHandler):
     def get(self):
         nbm = self.application.notebook_manager
         project = nbm.notebook_dir
         self.render('nbbrowser.html', project=project)
 
 
 class NewHandler(web.RequestHandler):
     def get(self):
         notebook_id = self.application.notebook_manager.new_notebook()
         self.render('notebook.html', notebook_id=notebook_id)
 
 
 class NamedNotebookHandler(web.RequestHandler):
     def get(self, notebook_id):
         nbm = self.application.notebook_manager
         if not nbm.notebook_exists(notebook_id):
             raise web.HTTPError(404)
         self.render('notebook.html', notebook_id=notebook_id)
 
 
 #-----------------------------------------------------------------------------
 # Kernel handlers
 #-----------------------------------------------------------------------------
 
 
 class MainKernelHandler(web.RequestHandler):
 
     def get(self):
         rkm = self.application.routing_kernel_manager
         self.finish(json.dumps(rkm.kernel_ids))
 
     def post(self):
         rkm = self.application.routing_kernel_manager
         notebook_id = self.get_argument('notebook', default=None)
         kernel_id = rkm.start_kernel(notebook_id)
         self.set_header('Location', '/'+kernel_id)
         self.finish(json.dumps(kernel_id))
 
 
 class KernelHandler(web.RequestHandler):
 
     SUPPORTED_METHODS = ('DELETE')
 
     def delete(self, kernel_id):
         rkm = self.application.routing_kernel_manager
         self.kill_kernel(kernel_id)
         self.set_status(204)
         self.finish()
 
 
 class KernelActionHandler(web.RequestHandler):
 
     def post(self, kernel_id, action):
         rkm = self.application.routing_kernel_manager
         if action == 'interrupt':
             rkm.interrupt_kernel(kernel_id)
             self.set_status(204)
         if action == 'restart':
             new_kernel_id = rkm.restart_kernel(kernel_id)
             self.write(json.dumps(new_kernel_id))
         self.finish()
 
 
 class ZMQStreamHandler(websocket.WebSocketHandler):
 
     def initialize(self, stream_name):
         self.stream_name = stream_name
 
     def open(self, kernel_id):
         rkm = self.application.routing_kernel_manager
         self.router = rkm.get_router(kernel_id, self.stream_name)
         self.client_id = self.router.register_client(self)
-        logging.info("Connection open: %s, %s" % (kernel_id, self.client_id))
 
     def on_message(self, msg):
         self.router.forward_msg(self.client_id, msg)
 
     def on_close(self):
         self.router.unregister_client(self.client_id)
-        logging.info("Connection closed: %s" % self.client_id)
 
 
 #-----------------------------------------------------------------------------
 # Notebook web service handlers
 #-----------------------------------------------------------------------------
 
 class NotebookRootHandler(web.RequestHandler):
 
     def get(self):
         nbm = self.application.notebook_manager
         files = nbm.list_notebooks()
         self.finish(json.dumps(files))
 
     def post(self):
         nbm = self.application.notebook_manager
         body = self.request.body.strip()
         format = self.get_argument('format', default='json')
         name = self.get_argument('name', default=None)
         if body:
             notebook_id = nbm.save_new_notebook(body, name=name, format=format)
         else:
             notebook_id = nbm.new_notebook()
         self.set_header('Location', '/'+notebook_id)
         self.finish(json.dumps(notebook_id))
 
 
 class NotebookHandler(web.RequestHandler):
 
     SUPPORTED_METHODS = ('GET', 'PUT', 'DELETE')
 
     def get(self, notebook_id):
         nbm = self.application.notebook_manager
         format = self.get_argument('format', default='json')
         last_mod, name, data = nbm.get_notebook(notebook_id, format)
         if format == u'json':
             self.set_header('Content-Type', 'application/json')
             self.set_header('Content-Disposition','attachment; filename=%s.json' % name)
         elif format == u'xml':
             self.set_header('Content-Type', 'application/xml')
             self.set_header('Content-Disposition','attachment; filename=%s.ipynb' % name)
         elif format == u'py':
             self.set_header('Content-Type', 'application/x-python')
             self.set_header('Content-Disposition','attachment; filename=%s.py' % name)
         self.set_header('Last-Modified', last_mod)
         self.finish(data)
 
     def put(self, notebook_id):
         nbm = self.application.notebook_manager
         format = self.get_argument('format', default='json')
         name = self.get_argument('name', default=None)
         nbm.save_notebook(notebook_id, self.request.body, name=name, format=format)
         self.set_status(204)
         self.finish()
 
     def delete(self, notebook_id):
         nbm = self.application.notebook_manager
         nbm.delete_notebook(notebook_id)
         self.set_status(204)
         self.finish()
 

IPython/frontend/html/notebook/kernelmanager.py

@@ -1,306 +1,350 b''
 """A kernel manager for multiple kernels."""
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING.txt, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import signal
 import sys
 import uuid
 
 import zmq
 
 from tornado import web
 
 from .routers import IOPubStreamRouter, ShellStreamRouter
 
 from IPython.config.configurable import LoggingConfigurable
 from IPython.zmq.ipkernel import launch_kernel
 from IPython.utils.traitlets import Instance, Dict, List, Unicode
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 class DuplicateKernelError(Exception):
     pass
 
 
 class KernelManager(LoggingConfigurable):
     """A class for managing multiple kernels."""
 
     context = Instance('zmq.Context')
     def _context_default(self):
         return zmq.Context.instance()
 
     _kernels = Dict()
 
     @property
     def kernel_ids(self):
         """Return a list of the kernel ids of the active kernels."""
         return self._kernels.keys()
 
     def __len__(self):
         """Return the number of running kernels."""
         return len(self.kernel_ids)
 
     def __contains__(self, kernel_id):
         if kernel_id in self.kernel_ids:
             return True
         else:
             return False
 
     def start_kernel(self, **kwargs):
         """Start a new kernel."""
         kernel_id = unicode(uuid.uuid4())
         (process, shell_port, iopub_port, stdin_port, hb_port) = launch_kernel(**kwargs)
         # Store the information for contacting the kernel. This assumes the kernel is
         # running on localhost.
         d = dict(
             process = process,
             stdin_port = stdin_port,
             iopub_port = iopub_port,
             shell_port = shell_port,
             hb_port = hb_port,
             ip = '127.0.0.1'
         )
         self._kernels[kernel_id] = d
         return kernel_id
 
     def kill_kernel(self, kernel_id):
         """Kill a kernel by its kernel uuid.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel to kill.
         """
         kernel_process = self.get_kernel_process(kernel_id)
         if kernel_process is not None:
             # Attempt to kill the kernel.
             try:
                 kernel_process.kill()
             except OSError, e:
                 # In Windows, we will get an Access Denied error if the process
                 # has already terminated. Ignore it.
                 if not (sys.platform == 'win32' and e.winerror == 5):
                     raise
             del self._kernels[kernel_id]
 
     def interrupt_kernel(self, kernel_id):
         """Interrupt (SIGINT) the kernel by its uuid.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel to interrupt.
         """
         kernel_process = self.get_kernel_process(kernel_id)
         if kernel_process is not None:
             if sys.platform == 'win32':
                 from parentpoller import ParentPollerWindows as Poller
                 Poller.send_interrupt(kernel_process.win32_interrupt_event)
             else:
                 kernel_process.send_signal(signal.SIGINT)
 
     def signal_kernel(self, kernel_id, signum):
         """ Sends a signal to the kernel by its uuid.
 
         Note that since only SIGTERM is supported on Windows, this function
         is only useful on Unix systems.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel to signal.
         """
         kernel_process = self.get_kernel_process(kernel_id)
         if kernel_process is not None:
             kernel_process.send_signal(signum)
 
     def get_kernel_process(self, kernel_id):
         """Get the process object for a kernel by its uuid.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel.
         """
         d = self._kernels.get(kernel_id)
         if d is not None:
             return d['process']
         else:
             raise KeyError("Kernel with id not found: %s" % kernel_id)
 
     def get_kernel_ports(self, kernel_id):
         """Return a dictionary of ports for a kernel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel.
 
         Returns
         =======
         port_dict : dict
             A dict of key, value pairs where the keys are the names
             (stdin_port,iopub_port,shell_port) and the values are the
             integer port numbers for those channels.
         """
         d = self._kernels.get(kernel_id)
         if d is not None:
             dcopy = d.copy()
             dcopy.pop('process')
             dcopy.pop('ip')
             return dcopy
         else:
             raise KeyError("Kernel with id not found: %s" % kernel_id)
 
     def get_kernel_ip(self, kernel_id):
         """Return ip address for a kernel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel.
 
         Returns
         =======
         ip : str
             The ip address of the kernel.
         """
         d = self._kernels.get(kernel_id)
         if d is not None:
             return d['ip']
         else:
             raise KeyError("Kernel with id not found: %s" % kernel_id)
 
     def create_session_manager(self, kernel_id):
         """Create a new session manager for a kernel by its uuid."""
         from sessionmanager import SessionManager
         return SessionManager(
             kernel_id=kernel_id, kernel_manager=self, 
             config=self.config, context=self.context, log=self.log
         )
 
 
 class RoutingKernelManager(LoggingConfigurable):
     """A KernelManager that handles WebSocket routing and HTTP error handling"""
 
     kernel_argv = List(Unicode)
     kernel_manager = Instance(KernelManager)
 
     _routers = Dict()
     _session_dict = Dict()
     _notebook_mapping = Dict()
 
     #-------------------------------------------------------------------------
     # Methods for managing kernels and sessions
     #-------------------------------------------------------------------------
 
     @property
     def kernel_ids(self):
+        """List the kernel ids."""
         return self.kernel_manager.kernel_ids
 
+    def kernel_for_notebook(self, notebook_id):
+        """Return the kernel_id for a notebook_id or None."""
+        return self._notebook_mapping.get(notebook_id)
+
+    def set_kernel_for_notebook(self, notebook_id, kernel_id):
+        """Associate a notebook with a kernel."""
+        if notebook_id is not None:
+            self._notebook_mapping[notebook_id] = kernel_id
+
     def notebook_for_kernel(self, kernel_id):
+        """Return the notebook_id for a kernel_id or None."""
         notebook_ids = [k for k, v in self._notebook_mapping.iteritems() if v == kernel_id]
         if len(notebook_ids) == 1:
             return notebook_ids[0]
         else:
             return None
 
     def delete_mapping_for_kernel(self, kernel_id):
+        """Remove the kernel/notebook mapping for kernel_id."""
         notebook_id = self.notebook_for_kernel(kernel_id)
         if notebook_id is not None:
             del self._notebook_mapping[notebook_id]
 
     def start_kernel(self, notebook_id=None):
+        """Start a kernel an return its kernel_id.
+
+        Parameters
+        ----------
+        notebook_id : uuid
+            The uuid of the notebook to associate the new kernel with. If this
+            is not None, this kernel will be persistent whenever the notebook
+            requests a kernel.
+        """
         self.log.info
-        kernel_id = self._notebook_mapping.get(notebook_id)
+        kernel_id = self.kernel_for_notebook(notebook_id)
         if kernel_id is None:
             kwargs = dict()
             kwargs['extra_arguments'] = self.kernel_argv
             kernel_id = self.kernel_manager.start_kernel(**kwargs)
-            if notebook_id is not None:
-                self._notebook_mapping[notebook_id] = kernel_id
-            self.log.info("Kernel started for notebook %s: %s" % (notebook_id,kernel_id))
+            self.set_kernel_for_notebook(notebook_id, kernel_id)
+            self.log.info("Kernel started: %s" % kernel_id)
             self.log.debug("Kernel args: %r" % kwargs)
             self.start_session_manager(kernel_id)
         else:
             self.log.info("Using existing kernel: %s" % kernel_id)
         return kernel_id
 
     def start_session_manager(self, kernel_id):
+        """Start the ZMQ sockets (a "session") to connect to a kernel."""
         sm = self.kernel_manager.create_session_manager(kernel_id)
         self._session_dict[kernel_id] = sm
         iopub_stream = sm.get_iopub_stream()
         shell_stream = sm.get_shell_stream()
         iopub_router = IOPubStreamRouter(
             zmq_stream=iopub_stream, session=sm.session, config=self.config
         )
         shell_router = ShellStreamRouter(
             zmq_stream=shell_stream, session=sm.session, config=self.config
         )
-        self._routers[(kernel_id, 'iopub')] = iopub_router
-        self._routers[(kernel_id, 'shell')] = shell_router
+        self.set_router(kernel_id, 'iopub', iopub_router)
+        self.set_router(kernel_id, 'shell', shell_router)
 
     def kill_kernel(self, kernel_id):
+        """Kill a kernel and remove its notebook association."""
         if kernel_id not in self.kernel_manager:
             raise web.HTTPError(404)
         try:
             sm = self._session_dict.pop(kernel_id)
         except KeyError:
             raise web.HTTPError(404)
         sm.stop()
         self.kernel_manager.kill_kernel(kernel_id)
         self.delete_mapping_for_kernel(kernel_id)
         self.log.info("Kernel killed: %s" % kernel_id)
 
     def interrupt_kernel(self, kernel_id):
+        """Interrupt a kernel."""
         if kernel_id not in self.kernel_manager:
             raise web.HTTPError(404)
         self.kernel_manager.interrupt_kernel(kernel_id)
         self.log.debug("Kernel interrupted: %s" % kernel_id)
 
     def restart_kernel(self, kernel_id):
+        """Restart a kernel while keeping clients connected."""
         if kernel_id not in self.kernel_manager:
             raise web.HTTPError(404)
 
         # Get the notebook_id to preserve the kernel/notebook association
         notebook_id = self.notebook_for_kernel(kernel_id)
         # Create the new kernel first so we can move the clients over.
         new_kernel_id = self.start_kernel()
 
         # Copy the clients over to the new routers.
         old_iopub_router = self.get_router(kernel_id, 'iopub')
         old_shell_router = self.get_router(kernel_id, 'shell')
         new_iopub_router = self.get_router(new_kernel_id, 'iopub')
         new_shell_router = self.get_router(new_kernel_id, 'shell')
         new_iopub_router.copy_clients(old_iopub_router)
         new_shell_router.copy_clients(old_shell_router)
 
+        # Shut down the old routers
+        old_shell_router.close()
+        old_iopub_router.close()
+        self.delete_router(kernel_id, 'shell')
+        self.delete_router(kernel_id, 'iopub')
+        del old_shell_router
+        del old_iopub_router
+
         # Now shutdown the old session and the kernel.
         # TODO: This causes a hard crash in ZMQStream.close, which sets
         # self.socket to None to hastily. We will need to fix this in PyZMQ
         # itself. For now, we just leave the old kernel running :(
         # Maybe this is fixed now, but nothing was changed really.
         self.kill_kernel(kernel_id)
 
         # Now save the new kernel/notebook association. We have to save it
         # after the old kernel is killed as that will delete the mapping.
-        self._notebook_mapping[notebook_id] = kernel_id
+        self.set_kernel_for_notebook(notebook_id, new_kernel_id)
 
-        self.log.debug("Kernel restarted: %s -> %s" % (kernel_id, new_kernel_id))
+        self.log.debug("Kernel restarted: %s" % new_kernel_id)
         return new_kernel_id
 
     def get_router(self, kernel_id, stream_name):
+        """Return the router for a given kernel_id and stream name."""
         router = self._routers[(kernel_id, stream_name)]
         return router
 
+    def set_router(self, kernel_id, stream_name, router):
+        """Set the router for a given kernel_id and stream_name."""
+        self._routers[(kernel_id, stream_name)] = router
+
+    def delete_router(self, kernel_id, stream_name):
+        """Delete a router for a kernel_id and stream_name."""
+        try:
+            del self._routers[(kernel_id, stream_name)]
+        except KeyError:
+            pass
+

IPython/frontend/html/notebook/routers.py

@@ -1,118 +1,125 b''
 """Routers that connect WebSockets to ZMQ sockets."""
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING.txt, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import uuid
 from Queue import Queue
 import json
 
 from IPython.config.configurable import Configurable
 from IPython.utils.traitlets import Instance, Int, Dict
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 class ZMQStreamRouter(Configurable):
 
     zmq_stream = Instance('zmq.eventloop.zmqstream.ZMQStream')
     session = Instance('IPython.zmq.session.Session')
     max_msg_size = Int(2048, config=True, help="""
         The max raw message size accepted from the browser
         over a WebSocket connection.
     """)
 
     _clients = Dict()
 
     def __init__(self, **kwargs):
         super(ZMQStreamRouter,self).__init__(**kwargs)
         self.zmq_stream.on_recv(self._on_zmq_reply)
 
+    def __del__(self):
+        self.close()
+
+    def close(self):
+        """Disable the routing actions of this router."""
+        self._clients = {}
+        self.zmq_stream.on_recv(None)
+
     def register_client(self, client):
         """Register a client, returning a client uuid."""
-        client_id = uuid.uuid4()
+        client_id = unicode(uuid.uuid4())
         self._clients[client_id] = client
         return client_id
 
     def unregister_client(self, client_id):
         """Unregister a client by its client uuid."""
         del self._clients[client_id]
 
     def copy_clients(self, router):
         """Copy the clients of another router to this one.
 
         This is used to enable the backend zeromq stream to disconnect
         and reconnect while the WebSocket connections to browsers
         remain, such as when a kernel is restarted.
         """
         for client_id, client in router._clients.items():
             client.router = self
             self._clients[client_id] = client
 
     def forward_msg(self, client_id, msg):
         """Forward a msg to a client by its id.
 
         The default implementation of this will fail silently if a message
         arrives on a socket that doesn't support it. This method should
         use max_msg_size to check and silently discard message that are too
         long."""
         pass
 
     def _on_zmq_reply(self, msg_list):
         """Handle a message the ZMQ stream sends to the router.
 
         Usually, this is where the return message will be written to
         clients that need it using client.write_message().
         """
         pass
 
     def _reserialize_reply(self, msg_list):
         """Reserialize a reply message using JSON.
 
         This takes the msg list from the ZMQ socket, unserializes it using
         self.session and then serializes the result using JSON. This method
         should be used by self._on_zmq_reply to build messages that can
         be sent back to the browser.
         """
         idents, msg_list = self.session.feed_identities(msg_list)
         msg = self.session.unserialize(msg_list)
         msg['header'].pop('date')
         msg.pop('buffers')
         return json.dumps(msg)
 
 
 class IOPubStreamRouter(ZMQStreamRouter):
 
     def _on_zmq_reply(self, msg_list):
         msg = self._reserialize_reply(msg_list)
         for client_id, client in self._clients.items():
             client.write_message(msg)
 
 
 class ShellStreamRouter(ZMQStreamRouter):
 
     _request_queue = Instance(Queue,(),{})
 
     def _on_zmq_reply(self, msg_list):
         msg = self._reserialize_reply(msg_list)
         client_id = self._request_queue.get(block=False)
         client = self._clients.get(client_id)
         if client is not None:
             client.write_message(msg)
 
     def forward_msg(self, client_id, msg):
         if len(msg) < self.max_msg_size:
             msg = json.loads(msg)
-            # to_send = self.session.serialize(msg)
             self._request_queue.put(client_id)        
             self.session.send(self.zmq_stream, msg)