Add joined IOPubAndShellHandler

minrk · minrk · commit bbfe2c7fbada · 2012-09-04T16:19:05.000-07:00
Now only one SockJS connection is needed in the browser for both channels.
The previous single-channel connections remain available.
diff --git a/IPython/frontend/html/notebook/handlers.py b/IPython/frontend/html/notebook/handlers.py
@@ -384,7 +384,7 @@ def application(self):
     def request(self):
         return self.session.handler.request
 
-    def _reserialize_reply(self, msg_list):
+    def _reserialize_reply(self, msg_list, channel=None):
         """Reserialize a reply message using JSON.
 
         This takes the msg list from the ZMQ socket, unserializes it using
@@ -394,20 +394,13 @@ def _reserialize_reply(self, msg_list):
         """
         idents, msg_list = self.ip_session.feed_identities(msg_list)
         msg = self.ip_session.unserialize(msg_list)
-        try:
-            msg['header'].pop('date')
-        except KeyError:
-            pass
-        try:
-            msg['parent_header'].pop('date')
-        except KeyError:
-            pass
-        msg.pop('buffers')
+        msg.pop('buffers', None)
+        msg['channel'] = channel
         return jsonapi.dumps(msg, default=date_default)
 
-    def _on_zmq_reply(self, msg_list):
+    def _on_zmq_reply(self, msg_list, channel=None):
         try:
-            msg = self._reserialize_reply(msg_list)
+            msg = self._reserialize_reply(msg_list, channel)
         except Exception:
             self.application.log.critical("Malformed message: %r" % msg_list, exc_info=True)
         else:
@@ -440,6 +433,12 @@ def on_open(self, info):
         self.ip_session = Session(config=cfg)
         self.save_on_message = self.on_message
         self.on_message = self.on_first_message
+        
+        self.initialize()
+    
+    def initialize(self):
+        """override in subclasses"""
+        pass
 
     def get_current_user(self):
         handler = self.session.handler
@@ -465,6 +464,11 @@ def on_first_message(self, cookie):
             logging.warn("Couldn't authenticate WebSocket connection")
             raise web.HTTPError(403)
         self.on_message = self.save_on_message
+        self.create_streams()
+    
+    def create_streams(self):
+        """override in subclass"""
+        pass
 
 
 class IOPubHandler(AuthenticatedZMQStreamHandler):
@@ -473,28 +477,23 @@ class IOPubHandler(AuthenticatedZMQStreamHandler):
     _beating = False
     iopub_stream = None
     hb_stream = None
-
-    def on_first_message(self, msg):
-        try:
-            super(IOPubHandler, self).on_first_message(msg)
-        except web.HTTPError:
-            self.close()
-            return
+    
+    def initialize(self):
         km = self.application.kernel_manager
         self.time_to_dead = km.time_to_dead
         self.first_beat = km.first_beat
+    
+    def create_streams(self):
+        km = self.application.kernel_manager
         kernel_id = self.kernel_id
         try:
             self.iopub_stream = km.create_iopub_stream(kernel_id)
             self.hb_stream = km.create_hb_stream(kernel_id)
-        except web.HTTPError:
-            # WebSockets don't response to traditional error codes so we
-            # close the connection.
-            if not self.stream.closed():
-                self.stream.close()
+        except Exception:
+            logging.error("Couldn't create streams", exc_info=True)
             self.close()
         else:
-            self.iopub_stream.on_recv(self._on_zmq_reply)
+            self.iopub_stream.on_recv(lambda msg: self._on_zmq_reply(msg, 'iopub'))
             self.start_hb(self.kernel_died)
 
     def on_message(self, msg):
@@ -562,36 +561,32 @@ def kernel_died(self):
         self.send(json.dumps(
             {'header': {'msg_type': 'status'},
              'parent_header': {},
-             'content': {'execution_state':'dead'}
+             'content': {'execution_state':'dead'},
+             'channel': 'iopub',
             }
         ))
         self.on_close()
 
 
 class ShellHandler(AuthenticatedZMQStreamHandler):
 
-    def initialize(self, *args, **kwargs):
-        self.shell_stream = None
-
-    def on_first_message(self, msg):
-        try:
-            super(ShellHandler, self).on_first_message(msg)
-        except web.HTTPError:
-            self.close()
-            return
+    shell_stream = None
+    max_msg_size = 65535
+    
+    def initialize(self):
         km = self.application.kernel_manager
         self.max_msg_size = km.max_msg_size
+    
+    def create_streams(self):
+        km = self.application.kernel_manager
         kernel_id = self.kernel_id
         try:
             self.shell_stream = km.create_shell_stream(kernel_id)
-        except web.HTTPError:
-            # WebSockets don't response to traditional error codes so we
-            # close the connection.
-            if not self.stream.closed():
-                self.stream.close()
+        except Exception:
+            logging.error("Couldn't create shell stream", exc_info=True)
             self.close()
         else:
-            self.shell_stream.on_recv(self._on_zmq_reply)
+            self.shell_stream.on_recv(lambda msg: self._on_zmq_reply(msg, 'shell'))
 
     def on_message(self, msg):
         if len(msg) < self.max_msg_size:
@@ -603,6 +598,20 @@ def on_close(self):
         if self.shell_stream is not None and not self.shell_stream.closed():
             self.shell_stream.close()
 
+class IOPubAndShellHandler(ShellHandler, IOPubHandler):
+    """single SockJS handler for IOPub + Shell zmq channels"""
+    def on_close(self):
+        ShellHandler.on_close(self)
+        IOPubHandler.on_close(self)
+    
+    def initialize(self):
+        ShellHandler.initialize(self)
+        IOPubHandler.initialize(self)
+    
+    def create_streams(self):
+        ShellHandler.create_streams(self)
+        IOPubHandler.create_streams(self)
+
 
 #-----------------------------------------------------------------------------
 # Notebook web service handlers
diff --git a/IPython/frontend/html/notebook/notebookapp.py b/IPython/frontend/html/notebook/notebookapp.py
@@ -52,7 +52,7 @@
     ShellHandler, NotebookRootHandler, NotebookHandler, NotebookCopyHandler,
     RSTHandler, AuthenticatedFileHandler, PrintNotebookHandler,
     MainClusterHandler, ClusterProfileHandler, ClusterActionHandler,
-    FileFindHandler,
+    FileFindHandler, IOPubAndShellHandler,
 )
 from .nbmanager import NotebookManager
 from .filenbmanager import FileNotebookManager
@@ -152,9 +152,11 @@ def __init__(self, ipython_app, kernel_manager, notebook_manager,
         
         iopub_router = SockJSRouter(IOPubHandler, r"/kernels/%s/iopub" % _uuid_regex)
         shell_router = SockJSRouter(ShellHandler, r"/kernels/%s/shell" % _uuid_regex)
+        sock_router = SockJSRouter(IOPubAndShellHandler, r"/kernels/%s/sock" % _uuid_regex)
         
         handlers += shell_router.urls
         handlers += iopub_router.urls
+        handlers += sock_router.urls
 
         # Python < 2.6.5 doesn't accept unicode keys in f(**kwargs), and
         # base_project_url will always be unicode, which will in turn
diff --git a/IPython/frontend/html/notebook/static/js/kernel.js b/IPython/frontend/html/notebook/static/js/kernel.js
@@ -18,7 +18,6 @@ var IPython = (function (IPython) {
     var Kernel = function (base_url) {
         this.kernel_id = null;
         this.shell_channel = null;
-        this.iopub_channel = null;
         this.base_url = base_url;
         this.running = false;
         this.username = "username";
@@ -71,13 +70,22 @@ var IPython = (function (IPython) {
 
     Kernel.prototype._kernel_started = function (json) {
         console.log("Kernel started: ", json.kernel_id);
+        var that = this
         this.running = true;
         this.kernel_id = json.kernel_id;
         this.ws_url = json.ws_url;
         this.kernel_url = this.base_url + "/" + this.kernel_id;
         this.start_channels();
-        this.shell_channel.onmessage = $.proxy(this._handle_shell_reply,this);
-        this.iopub_channel.onmessage = $.proxy(this._handle_iopub_reply,this);
+        this.shell_channel.onmessage = function(e) {
+            var msg = $.parseJSON(e.data);
+            if (msg['channel'] == 'shell'){
+                that._handle_shell_reply(e);
+            } else if (msg['channel'] == 'iopub') {
+                that._handle_iopub_reply(e);
+            } else {
+                console.log("Bad message", e);
+            }
+        };
     };
 
 
@@ -118,8 +126,7 @@ var IPython = (function (IPython) {
         this.stop_channels();
         var ws_url = this.ws_url + this.kernel_url;
         console.log("Starting SockJS:", ws_url);
-        this.shell_channel = new SockJS(ws_url + "/shell");
-        this.iopub_channel = new SockJS(ws_url + "/iopub");
+        this.shell_channel = new SockJS(ws_url + "/sock");
         send_cookie = function(){
             this.send(document.cookie);
         };
@@ -144,12 +151,9 @@ var IPython = (function (IPython) {
         };
         this.shell_channel.onopen = send_cookie;
         this.shell_channel.onclose = ws_closed_early;
-        this.iopub_channel.onopen = send_cookie;
-        this.iopub_channel.onclose = ws_closed_early;
         // switch from early-close to late-close message after 1s
         setTimeout(function(){
             that.shell_channel.onclose = ws_closed_late;
-            that.iopub_channel.onclose = ws_closed_late;
         }, 1000);
     };
 
@@ -160,11 +164,6 @@ var IPython = (function (IPython) {
             this.shell_channel.close();
             this.shell_channel = null;
         };
-        if (this.iopub_channel !== null) {
-            this.iopub_channel.onclose = function (evt) {};
-            this.iopub_channel.close();
-            this.iopub_channel = null;
-        };
     };
 
     // Main public methods.
