Switch to EventSource

nbresuse/__init__.py

@@ -1,37 +1,41 @@
 import os
 import json
-import psutil
-from traitlets import Float, Int, default
-from traitlets.config import Configurable
 from notebook.utils import url_path_join
 from notebook.base.handlers import IPythonHandler
-from tornado import web
+from tornado import web, iostream
+import asyncio
+import pluggy
+from nbresuse import hooks, default_resources
+from collections import ChainMap
+
+
+plugin_manager = pluggy.PluginManager('nbresuse')
+plugin_manager.add_hookspecs(hooks)
+# Register the resources nbresuse provides by default
+plugin_manager.register(default_resources)
 
 
 class MetricsHandler(IPythonHandler):
+    def initialize(self, nbapp):
+        self.set_header('content-type', 'text/event-stream')
+        self.set_header('cache-control', 'no-cache')
+        self.nbapp = nbapp
+
     @web.authenticated
-    def get(self):
+    async def get(self):
         """
         Calculate and return current resource usage metrics
         """
-        config = self.settings['nbresuse_display_config']
-        cur_process = psutil.Process()
-        all_processes = [cur_process] + cur_process.children(recursive=True)
-        rss = sum([p.memory_info().rss for p in all_processes])
-
-        limits = {}
-
-        if config.mem_limit != 0:
-            limits['memory'] = {
-                'rss': config.mem_limit
-            }
-            if config.mem_warning_threshold != 0:
-                limits['memory']['warn'] = (config.mem_limit - rss) < (config.mem_limit * config.mem_warning_threshold)
-        metrics = {
-            'rss': rss,
-            'limits': limits,
-        }
-        self.write(json.dumps(metrics))
+        while True:
+            metrics = {}
+            for metric_response in plugin_manager.hook.nbresuse_add_resource(nbapp=self.nbapp):
+                metrics.update(metric_response)
+            self.write('data: {}\n\n'.format(json.dumps(metrics)))
+            try:
+                await self.flush()
+            except iostream.StreamClosedError:
+                return
+            await asyncio.sleep(5)
 
 
 def _jupyter_server_extension_paths():
@@ -53,47 +57,11 @@ def _jupyter_nbextension_paths():
         "require": "nbresuse/main"
     }]
 
-class ResourceUseDisplay(Configurable):
-    """
-    Holds server-side configuration for nbresuse
-    """
-
-    mem_warning_threshold = Float(
-        0.1,
-        help="""
-        Warn user with flashing lights when memory usage is within this fraction
-        memory limit.
-
-        For example, if memory limit is 128MB, `mem_warning_threshold` is 0.1,
-        we will start warning the user when they use (128 - (128 * 0.1)) MB.
-
-        Set to 0 to disable warning.
-        """,
-        config=True
-    )
-
-    mem_limit = Int(
-        0,
-        config=True,
-        help="""
-        Memory limit to display to the user, in bytes.
-
-        Note that this does not actually limit the user's memory usage!
-
-        Defaults to reading from the `MEM_LIMIT` environment variable. If
-        set to 0, no memory limit is displayed.
-        """
-    )
-
-    @default('mem_limit')
-    def _mem_limit_default(self):
-        return int(os.environ.get('MEM_LIMIT', 0))
-
 def load_jupyter_server_extension(nbapp):
     """
     Called during notebook start
     """
-    resuseconfig = ResourceUseDisplay(parent=nbapp)
-    nbapp.web_app.settings['nbresuse_display_config'] = resuseconfig
-    route_pattern = url_path_join(nbapp.web_app.settings['base_url'], '/metrics')
-    nbapp.web_app.add_handlers('.*', [(route_pattern, MetricsHandler)])
+    api_url= url_path_join(nbapp.web_app.settings['base_url'], '/api/nbresuse/v1')
+    nbapp.web_app.add_handlers('.*', [
+        (api_url, MetricsHandler, {'nbapp': nbapp})
+    ])

nbresuse/default_resources.py

@@ -0,0 +1,48 @@
+from nbresuse.hooks import hookimpl
+import psutil
+import os
+from traitlets.config import Configurable
+from traitlets import Float, List
+
+
+class ResourceUseDisplay(Configurable):
+    """
+    Holds server-side configuration for nbresuse
+    """
+
+    mem_warning_threshold = Float(
+        0.9,
+        help="""
+        Warn user with flashing lights when memory usage is
+        more than this fraction of memory limit.
+
+        For example, if memory limit is 128MB,
+        `mem_warning_threshold` is 0.9,
+        we will start warning the user when they use
+        (128 * 0.9) MB.
+
+        Set to 0 to disable warning.
+        """,
+        config=True
+    )
+
+
+# NOTE: This hits /proc so many times it is gonna make *someone* mad
+# FIXME: Cache this!
+@hookimpl
+def nbresuse_add_resource(nbapp):
+    cur_process = psutil.Process()
+    all_processes = [cur_process] + cur_process.children(recursive=True)
+
+    config = ResourceUseDisplay(parent=nbapp)
+    rss = sum([p.memory_info().rss for p in all_processes])
+
+    return {
+        'nbresuse.jupyter.org/usage': {
+            'rss':{
+                'usage': rss,
+                'limit': os.environ.get('MEM_LIMIT', None),
+                'warn': config.mem_warning_threshold
+            }
+        }
+    }

nbresuse/hooks.py

@@ -0,0 +1,14 @@
+import pluggy
+
+hookspec = pluggy.HookspecMarker("nbresuse")
+hookimpl = pluggy.HookimplMarker("nbresuse")
+
+@hookspec
+def nbresuse_add_resource(nbapp):
+    """
+    Return resource definitions to send to clients.
+
+    Should return a dictionary that'll be merged with all other
+    resources to be sent to the client.
+    """
+

nbresuse/static/main.js

@@ -25,41 +25,37 @@ define(['jquery', 'base/js/utils'], function ($, utils) {
             // Don't poll when nobody is looking
             return;
         }
-        $.getJSON(utils.get_body_data('baseUrl') + 'metrics', function(data) {
+        var metricsUrl = utils.get_body_data('baseUrl') + 'api/nbresuse/v1';
+
+        // FIXME: Reconnect on failure
+        var metricsSource = new EventSource(metricsUrl);
+        metricsSource.onmessage = function(message) {
+            var data = JSON.parse(message.data);
             // FIXME: Proper setups for MB and GB. MB should have 0 things
             // after the ., but GB should have 2.
-            var display = Math.round(data['rss'] / (1024 * 1024));
 
-            var limits = data['limits'];
-            if ('memory' in limits) {
-                if ('rss' in limits['memory']) {
-                    display += " / " + (limits['memory']['rss'] / (1024 * 1024));
-                }
-                if (limits['memory']['warn']) {
+            var memory = data['nbresuse.jupyter.org/usage'];
+
+            // Show RSS info
+            var rss = memory['rss'];
+            var display = Math.round(rss['usage'] / (1024 * 1024));
+
+            if (rss['limit']) {
+                display += ' / ' + Math.round((rss['limit'] / (1024 * 1024)));
+
+                if (rss['usage'] / rss['limit'] >= rss['warn']) {
                     $('#nbresuse-display').addClass('nbresuse-warn');
                 } else {
                     $('#nbresuse-display').removeClass('nbresuse-warn');
                 }
             }
-            if (data['limits']['memory'] !== null) {
-            }
             $('#nbresuse-mem').text(display + ' MB');
-        });
-    }
+        };
+    };
 
     var load_ipython_extension = function () {
         setupDOM();
         displayMetrics();
-        // Update every five seconds, eh?
-        setInterval(displayMetrics, 1000 * 5);
-
-        document.addEventListener("visibilitychange", function() {
-            // Update instantly when user activates notebook tab
-            // FIXME: Turn off update timer completely when tab not in focus
-            if (!document.hidden) {
-                displayMetrics();
-            }
-        }, false);
     };
 
     return {

setup.py

@@ -11,6 +11,7 @@
     install_requires=[
         'psutil',
         'notebook',
+        'pluggy'
     ],
     data_files=[
         ('share/jupyter/nbextensions/nbresuse', glob('nbresuse/static/*')),