More work on updating the notebook zmq forwarding.

IPython/frontend/html/notebook/handlers.py

@@ -55,7 +55,7 @@ def open(self, kernel_id):
         logging.info("Connection open: %s, %s" % (kernel_id, self.client_id))
 
     def on_message(self, msg):
-        self.router.forward_unicode(self.client_id, msg)
+        self.router.forward_msg(self.client_id, msg)
 
     def on_close(self):
         self.router.unregister_client(self.client_id)

IPython/frontend/html/notebook/notebookapp.py

@@ -55,7 +55,7 @@
 
 class NotebookWebApplication(web.Application):
 
-    def __init__(self, kernel_manager, log, kernel_argv):
+    def __init__(self, kernel_manager, log, kernel_argv, config):
         handlers = [
             (r"/", MainHandler),
             (r"/kernels", KernelHandler),
@@ -74,6 +74,7 @@ def __init__(self, kernel_manager, log, kernel_argv):
         self.kernel_manager = kernel_manager
         self.log = log
         self.kernel_argv = kernel_argv
+        self.config = config
         self._routers = {}
         self._session_dict = {}
 
@@ -99,8 +100,12 @@ def start_session_manager(self, kernel_id):
         self._session_dict[kernel_id] = sm
         iopub_stream = sm.get_iopub_stream()
         shell_stream = sm.get_shell_stream()
-        iopub_router = IOPubStreamRouter(iopub_stream, sm.session)
-        shell_router = ShellStreamRouter(shell_stream, sm.session)
+        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
 
@@ -230,7 +235,9 @@ def init_logging(self):
     def initialize(self, argv=None):
         super(IPythonNotebookApp, self).initialize(argv)
         self.init_kernel_manager()
-        self.web_app = NotebookWebApplication(self.kernel_manager, self.log, self.kernel_argv)
+        self.web_app = NotebookWebApplication(
+            self.kernel_manager, self.log, self.kernel_argv, self.config
+        )
         self.http_server = httpserver.HTTPServer(self.web_app)
         self.http_server.listen(self.port)
 

IPython/frontend/html/notebook/routers.py

@@ -2,57 +2,93 @@
 from Queue import Queue
 import json
 
-class ZMQStreamRouter(object):
+from IPython.config.configurable import Configurable
+from IPython.utils.traitlets import Instance, Int, Dict
 
-    def __init__(self, zmq_stream, session):
-        self.zmq_stream = zmq_stream
-        self.session = session
-        self._clients = {}
+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 register_client(self, client):
+        """Register a client, returning a client uuid."""
         client_id = 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.
+        """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
+
 
 class IOPubStreamRouter(ZMQStreamRouter):
 
     def _on_zmq_reply(self, msg_list):
+        msg = self.session.unpack_message(msg_list)
+        msg = json.dumps(msg)
         for client_id, client in self._clients.items():
             for msg in msg_list:
-                print "Got message: ", msg
                 client.write_message(msg)
 
-    def forward_unicode(self, client_id, msg):
-        # This is a SUB stream that we should never write to.
-        pass
-
 
 class ShellStreamRouter(ZMQStreamRouter):
 
-    def __init__(self, zmq_stream, session):
-        ZMQStreamRouter.__init__(self, zmq_stream, session)
-        self._request_queue = Queue()
+    _request_queue = Instance(Queue,(),{})
 
     def _on_zmq_reply(self, msg_list):
+        msg = self.session.unpack_message(msg_list)
+        msg = json.dumps(msg)
+        print "Reply: ", msg_list
         client_id = self._request_queue.get(block=False)
         client = self._clients.get(client_id)
         if client is not None:
             for msg in msg_list:
                 client.write_message(msg)
 
-    def forward_unicode(self, client_id, msg):
-        print "Inbound message: ", msg
-        self._request_queue.put(client_id)        
-        self.session.send(self.zmq_stream, msg)
+    def forward_msg(self, client_id, msg):
+        if len(msg) < self.max_msg_size:
+            msg = json.loads(msg)
+            print "Raw msg: ", msg
+            to_send = self.session.serialize(msg)
+            print "to_send: ", to_send, to_send[-3:]
+            self._request_queue.put(client_id)        
+            self.session.send_raw(self.zmq_stream, to_send[-3:])