track active kernel connections

notebook/services/kernels/handlers.py

@@ -260,6 +260,7 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
     
     def open(self, kernel_id):
         super(ZMQChannelsHandler, self).open()
+        self.kernel_manager.notify_connect(kernel_id)
         try:
             self.create_stream()
         except web.HTTPError as e:
@@ -401,6 +402,7 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
             self._open_sessions.pop(self.session_key)
         km = self.kernel_manager
         if self.kernel_id in km:
+            km.notify_disconnect(self.kernel_id)
             km.remove_restart_callback(
                 self.kernel_id, self.on_kernel_restarted,
             )

notebook/services/kernels/kernelmanager.py

@@ -14,7 +14,7 @@ from tornado.concurrent import Future
 from tornado.ioloop import IOLoop
 
 from jupyter_client.multikernelmanager import MultiKernelManager
-from traitlets import List, Unicode, TraitError, default, validate
+from traitlets import Dict, List, Unicode, TraitError, default, validate
 
 from notebook.utils import to_os_path
 from notebook.services.contents.tz import utcnow
@@ -31,6 +31,8 @@ class MappingKernelManager(MultiKernelManager):
     kernel_argv = List(Unicode())
 
     root_dir = Unicode(config=True)
+    
+    _kernel_connections = Dict()
 
     @default('root_dir')
     def _default_root_dir(self):
@@ -91,6 +93,7 @@ class MappingKernelManager(MultiKernelManager):
             kernel_id = yield gen.maybe_future(
                 super(MappingKernelManager, self).start_kernel(**kwargs)
             )
+            self._kernel_connections[kernel_id] = 0
             self.start_watching_activity(kernel_id)
             self.log.info("Kernel started: %s" % kernel_id)
             self.log.debug("Kernel args: %r" % kwargs)
@@ -109,6 +112,7 @@ class MappingKernelManager(MultiKernelManager):
         """Shutdown a kernel by kernel_id"""
         self._check_kernel_id(kernel_id)
         self._kernels[kernel_id]._activity_stream.close()
+        self._kernel_connections.pop(kernel_id, None)
         return super(MappingKernelManager, self).shutdown_kernel(kernel_id, now=now)
 
     def restart_kernel(self, kernel_id):
@@ -152,6 +156,16 @@ class MappingKernelManager(MultiKernelManager):
         timeout = loop.add_timeout(loop.time() + 30, on_timeout)
         return future
 
+    def notify_connect(self, kernel_id):
+        """Notice a new connection to a kernel"""
+        if kernel_id in self._kernel_connections:
+            self._kernel_connections[kernel_id] += 1
+
+    def notify_disconnect(self, kernel_id):
+        """Notice a disconnection from a kernel"""
+        if kernel_id in self._kernel_connections:
+            self._kernel_connections[kernel_id] -= 1
+
     def kernel_model(self, kernel_id):
         """Return a JSON-safe dict representing a kernel
 
@@ -165,6 +179,7 @@ class MappingKernelManager(MultiKernelManager):
             "name": kernel.kernel_name,
             "last_activity": kernel.last_activity,
             "execution_state": kernel.execution_state,
+            "connections": self._kernel_connections[kernel_id],
         }
         return model
 
@@ -204,7 +219,7 @@ class MappingKernelManager(MultiKernelManager):
             idents, fed_msg_list = kernel.session.feed_identities(msg_list)
             msg = kernel.session.deserialize(fed_msg_list)
             msg_type = msg['header']['msg_type']
-            self.log.info("activity on %s: %s", kernel_id, msg_type)
+            self.log.debug("activity on %s: %s", kernel_id, msg_type)
             if msg_type == 'status':
                 kernel.execution_state = msg['content']['execution_state']
 

notebook/services/kernels/tests/test_kernels_api.py

@@ -1,7 +1,10 @@
 """Test the kernels service API."""
 
 import json
+
 import requests
+from tornado.websocket import websocket_connect
+from tornado.ioloop import IOLoop
 
 from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
 
@@ -45,6 +48,14 @@ class KernelAPI(object):
     def restart(self, id):
         return self._req('POST', url_path_join(id, 'restart'))
 
+    def websocket(self, id):
+        loop = IOLoop()
+        f = websocket_connect(url_path_join(
+            self.base_url.replace('http', 'ws', 1), 'api/kernels', id, 'channels'),
+            io_loop=loop)
+        return loop.run_sync(lambda : f)
+
+
 class KernelAPITest(NotebookTestBase):
     """Test the kernels web service API"""
     def setUp(self):
@@ -144,3 +155,15 @@ class KernelAPITest(NotebookTestBase):
         bad_id = '111-111-111-111-111'
         with assert_http_error(404, 'Kernel does not exist: ' + bad_id):
             self.kern_api.shutdown(bad_id)
+
+    def test_connections(self):
+        kid = self.kern_api.start().json()['id']
+        model = self.kern_api.get(kid).json()
+        assert model['connections'] == 0
+
+        ws = self.kern_api.websocket(kid)
+        model = self.kern_api.get(kid).json()
+        assert model['connections'] == 1
+        ws.close()
+        model = self.kern_api.get(kid).json()
+        assert model['connections'] == 0