Merge pull request #23 from hadim/better-blocking

Run archive_and_download in an executor to make it really not blocking.

Author: hadim

jupyter_archive/handlers.py

@@ -9,6 +9,11 @@
 from notebook.utils import url2path
 
 
+# The delay in ms at which we send the chunk of data
+# to the client.
+ARCHIVE_DOWNLOAD_FLUSH_DELAY = 100
+
+
 class ArchiveStream():
     def __init__(self, handler):
         self.handler = handler
@@ -23,7 +28,10 @@ def tell(self):
         return self.position
 
     def flush(self):
-        self.handler.flush()
+        # Note: Flushing is done elsewhere, in the main thread
+        # because `write()` is called in a background thread.
+        # self.handler.flush()
+        pass
 
 
 def make_writer(handler, archive_format="zip"):
@@ -89,30 +97,40 @@ def get(self, archive_path, include_body=False):
         self.set_header('content-disposition',
                         'attachment; filename={}'.format(archive_filename))
 
-        task = asyncio.ensure_future(self.archive_and_download(archive_path, archive_format, archive_token))
+        self.canceled = False
+        self.flush_cb = ioloop.PeriodicCallback(self.flush, ARCHIVE_DOWNLOAD_FLUSH_DELAY)
+        self.flush_cb.start()
+
+        args = (archive_path, archive_format, archive_token)
+        yield ioloop.IOLoop.current().run_in_executor(None, self.archive_and_download, *args)
 
-        try:
-            yield task
-        except (asyncio.CancelledError, iostream.StreamClosedError):
-            task.cancel()
+        if self.canceled:
             self.log.info('Download canceled.')
         else:
+            self.flush()
             self.log.info('Finished downloading {}.'.format(archive_filename))
 
-    @gen.coroutine
+        self.set_cookie("archiveToken", archive_token)
+        self.flush_cb.stop()
+        self.finish()
+
     def archive_and_download(self, archive_path, archive_format, archive_token):
 
         with make_writer(self, archive_format) as archive:
             prefix = len(str(archive_path.parent)) + len(os.path.sep)
             for root, _, files in os.walk(archive_path):
                 for file_ in files:
                     file_name = os.path.join(root, file_)
-                    self.log.debug("{}\n".format(file_name))
-                    archive.add(file_name, os.path.join(root[prefix:], file_))
-                    yield self.flush()
-
-        self.set_cookie("archiveToken", archive_token)
-        self.finish()
+                    if not self.canceled:
+                        self.log.debug("{}\n".format(file_name))
+                        archive.add(file_name, os.path.join(root[prefix:], file_))
+                    else:
+                        break
+
+    def on_connection_close(self):
+        super().on_connection_close()
+        self.canceled = True
+        self.flush_cb.stop()
 
 
 class ExtractArchiveHandler(IPythonHandler):