iopub rate limit refinements

- reset counter on status: idle (avoids Run All running into limits)
- fix rate measurements to exclude messages not sent due to rate limiting
- fix formatting of rate-exceeded messages, tweak text a bit
- after pausing output, don't resume sending until rate has dropped 20% below limit, to avoid rapid back & forth across the limit

notebook/notebookapp.py

@@ -915,7 +915,7 @@ class NotebookApp(JupyterApp):
         help="Reraise exceptions encountered loading server extensions?",
     )
 
-    iopub_msg_rate_limit = Float(1000, config=True, help="""(msg/sec)
+    iopub_msg_rate_limit = Float(1000, config=True, help="""(msgs/sec)
         Maximum rate at which messages can be sent on iopub before they are
         limited.""")
 

notebook/services/kernels/handlers.py

@@ -8,6 +8,8 @@ Preliminary documentation at https://github.com/ipython/ipython/wiki/IPEP-16%3A-
 
 import json
 import logging
+from textwrap import dedent
+
 from tornado import gen, web
 from tornado.concurrent import Future
 from tornado.ioloop import IOLoop
@@ -108,11 +110,11 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
 
     @property
     def iopub_msg_rate_limit(self):
-        return self.settings.get('iopub_msg_rate_limit', None)
+        return self.settings.get('iopub_msg_rate_limit', 0)
 
     @property
     def iopub_data_rate_limit(self):
-        return self.settings.get('iopub_data_rate_limit', None)
+        return self.settings.get('iopub_data_rate_limit', 0)
 
     @property
     def rate_limit_window(self):
@@ -298,7 +300,7 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
         def write_stderr(error_message):
             self.log.warning(error_message)
             msg = self.session.msg("stream",
-                content={"text": error_message, "name": "stderr"},
+                content={"text": error_message + '\n', "name": "stderr"},
                 parent=parent
             )
             msg['channel'] = 'iopub'
@@ -306,6 +308,16 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
             
         channel = getattr(stream, 'channel', None)
         msg_type = msg['header']['msg_type']
+
+        if channel == 'iopub' and msg_type == 'status' and msg['content'].get('execution_state') == 'idle':
+            # reset rate limit counter on status=idle,
+            # to avoid 'Run All' hitting limits prematurely.
+            self._iopub_window_byte_queue = []
+            self._iopub_window_msg_count = 0
+            self._iopub_window_byte_count = 0
+            self._iopub_msgs_exceeded = False
+            self._iopub_data_exceeded = False
+
         if channel == 'iopub' and msg_type not in {'status', 'comm_open', 'execute_input'}:
             
             # Remove the counts queued for removal.
@@ -320,7 +332,7 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
                     # This part of the queue hasn't be reached yet, so we can
                     # abort the loop.
                     break
-                    
+
             # Increment the bytes and message count
             self._iopub_window_msg_count += 1
             byte_count = sum([len(x) for x in msg_list])
@@ -336,39 +348,45 @@ class ZMQChannelsHandler(AuthenticatedZMQStreamHandler):
             data_rate = float(self._iopub_window_byte_count) / self.rate_limit_window
             
             # Check the msg rate
-            if self.iopub_msg_rate_limit is not None and msg_rate > self.iopub_msg_rate_limit and self.iopub_msg_rate_limit > 0:
+            if self.iopub_msg_rate_limit > 0 and msg_rate > self.iopub_msg_rate_limit:
                 if not self._iopub_msgs_exceeded:
                     self._iopub_msgs_exceeded = True
-                    write_stderr("""iopub message rate exceeded.  The
-                    notebook server will temporarily stop sending iopub
-                    messages to the client in order to avoid crashing it.
+                    write_stderr(dedent("""\
+                    IOPub message rate exceeded.
+                    The notebook server will temporarily stop sending output
+                    to the client in order to avoid crashing it.
                     To change this limit, set the config variable
-                    `--NotebookApp.iopub_msg_rate_limit`.""")
-                    return
+                    `--NotebookApp.iopub_msg_rate_limit`."""))
             else:
-                if self._iopub_msgs_exceeded:
+                # resume once we've got some headroom below the limit
+                if self._iopub_msgs_exceeded and msg_rate < (0.8 * self.iopub_msg_rate_limit):
                     self._iopub_msgs_exceeded = False
                     if not self._iopub_data_exceeded:
                         self.log.warning("iopub messages resumed")
-    
+
             # Check the data rate
-            if self.iopub_data_rate_limit is not None and data_rate > self.iopub_data_rate_limit and self.iopub_data_rate_limit > 0:
+            if self.iopub_data_rate_limit > 0 and data_rate > self.iopub_data_rate_limit:
                 if not self._iopub_data_exceeded:
                     self._iopub_data_exceeded = True
-                    write_stderr("""iopub data rate exceeded.  The
-                    notebook server will temporarily stop sending iopub
-                    messages to the client in order to avoid crashing it.
+                    write_stderr(dedent("""\
+                    IOPub data rate exceeded.
+                    The notebook server will temporarily stop sending output
+                    to the client in order to avoid crashing it.
                     To change this limit, set the config variable
-                    `--NotebookApp.iopub_data_rate_limit`.""")
-                    return
+                    `--NotebookApp.iopub_data_rate_limit`."""))
             else:
-                if self._iopub_data_exceeded:
+                # resume once we've got some headroom below the limit
+                if self._iopub_data_exceeded and data_rate < (0.8 * self.iopub_data_rate_limit):
                     self._iopub_data_exceeded = False
                     if not self._iopub_msgs_exceeded:
                         self.log.warning("iopub messages resumed")
         
             # If either of the limit flags are set, do not send the message.
             if self._iopub_msgs_exceeded or self._iopub_data_exceeded:
+                # we didn't send it, remove the current message from the calculus
+                self._iopub_window_msg_count -= 1
+                self._iopub_window_byte_count -= byte_count
+                self._iopub_window_byte_queue.pop(-1)
                 return
         super(ZMQChannelsHandler, self)._on_zmq_reply(stream, msg)