fix memory leak #59 (#62)

* fix memory leak #59 by limiting both _write_buffer of RequstHandler, as well as _write_buffer of IOStream

* fix memory leak #59 set config=True and timeout for while loop

* Add configuration documentation

Co-authored-by: Frédéric Collonval <[email protected]>

Author: eastonsuo

README.md

@@ -20,6 +20,27 @@ Features:
 
 ![jupyter-archive in action](https://raw.githubusercontent.com/jupyterlab-contrib/jupyter-archive/master/archive.gif)
 
+## Configuration
+
+The server extension has some [configuration settings](https://jupyter-server.readthedocs.io/en/latest/users/configuration.html) -- 
+ the values below are the default one:
+
+```json5
+{
+  "JupyterArchive": {
+    "stream_max_buffer_size": 104857600, // The max size of tornado IOStream buffer
+    "handler_max_buffer_length": 10240, // The max length of chunks in tornado RequestHandler
+    "archive_download_flush_delay": 100 // The delay in ms at which we send the chunk of data to the client.
+  }
+}
+```
+
+You can also set new values with the following environment variables:
+
+- `JA_IOSTREAM_MAX_BUFFER_SIZE`
+- `JA_HANDLER_MAX_BUFFER_LENGTH`
+- `JA_ARCHIVE_DOWNLOAD_FLUSH_DELAY`
+
 ## Requirements
 
 - JupyterLab >= 3.0

jupyter_archive/__init__.py

@@ -1,5 +1,9 @@
 import json
 import pathlib
+import os
+
+from traitlets.config import Configurable
+from traitlets import Int, default
 
 from ._version import __version__
 from .handlers import setup_handlers
@@ -14,6 +18,31 @@ def _jupyter_labextension_paths():
     return [{"src": "labextension", "dest": data["name"]}]
 
 
+class JupyterArchive(Configurable):
+    stream_max_buffer_size = Int(help="The max size of tornado IOStream buffer",
+                                 config=True)
+
+    @default("stream_max_buffer_size")
+    def _default_stream_max_buffer_size(self):
+        # 100 * 1024 * 1024 equals to 100M
+        return int(os.environ.get("JA_IOSTREAM_MAX_BUFFER_SIZE", 100 * 1024 * 1024))
+
+    handler_max_buffer_length = Int(help="The max length of chunks in tornado RequestHandler",
+                                    config=True)
+
+    @default("handler_max_buffer_length")
+    def _default_handler_max_buffer_length(self):
+        # if 8K for one chunk, 10240 * 8K equals to 80M
+        return int(os.environ.get("JA_HANDLER_MAX_BUFFER_LENGTH", 10240))
+
+    archive_download_flush_delay = Int(help="The delay in ms at which we send the chunk of data to the client.",
+                                       config=True)
+
+    @default("archive_download_flush_delay")
+    def _default_archive_download_flush_delay(self):
+        return int(os.environ.get("JA_ARCHIVE_DOWNLOAD_FLUSH_DELAY", 100))
+
+
 def _jupyter_server_extension_points():
     return [{"module": "jupyter_archive"}]
 
@@ -26,6 +55,8 @@ def _load_jupyter_server_extension(server_app):
     server_app: jupyterlab.labapp.LabApp
         JupyterLab application instance
     """
+    config = JupyterArchive(config=server_app.config)
+    server_app.web_app.settings["jupyter_archive"] = config
     setup_handlers(server_app.web_app)
 
 

jupyter_archive/handlers.py

@@ -1,15 +1,13 @@
 import os
 import pathlib
 import tarfile
+import time
 import zipfile
 
 from jupyter_server.base.handlers import JupyterHandler
 from jupyter_server.utils import url2path, url_path_join
 from tornado import ioloop, web
 
-# The delay in ms at which we send the chunk of data
-# to the client.
-ARCHIVE_DOWNLOAD_FLUSH_DELAY = 100
 SUPPORTED_FORMAT = [
     "zip",
     "tgz",
@@ -29,6 +27,18 @@ def __init__(self, handler):
         self.position = 0
 
     def write(self, data):
+        if self.handler.canceled:
+            raise ValueError("File download canceled")
+        # timeout 600s for this while loop
+        time_out_cnt = 600 * 1000 / self.handler.archive_download_flush_delay
+        while len(self.handler._write_buffer) > self.handler.handler_max_buffer_length:
+            # write_buffer or handler is too large, wait for an flush cycle
+            time.sleep(self.handler.archive_download_flush_delay / 1000)
+            if self.handler.canceled:
+                raise ValueError("File download canceled")
+            time_out_cnt -= 1
+            if time_out_cnt <= 0:
+                raise ValueError("Time out for writing into tornado buffer")
         self.position += len(data)
         self.handler.write(data)
         del data
@@ -78,6 +88,26 @@ def make_reader(archive_path):
 
 
 class DownloadArchiveHandler(JupyterHandler):
+
+    @property
+    def stream_max_buffer_size(self):
+        return self.settings["jupyter_archive"].stream_max_buffer_size
+
+    @property
+    def handler_max_buffer_length(self):
+        return self.settings["jupyter_archive"].handler_max_buffer_length
+
+    @property
+    def archive_download_flush_delay(self):
+        return self.settings["jupyter_archive"].archive_download_flush_delay
+
+    def flush(self, include_footers=False):
+        # skip flush when stream_buffer is larger than stream_max_buffer_size
+        stream_buffer = self.request.connection.stream._write_buffer
+        if stream_buffer and len(stream_buffer) > self.stream_max_buffer_size:
+            return
+        return super(DownloadArchiveHandler, self).flush(include_footers)
+
     @web.authenticated
     async def get(self, archive_path, include_body=False):
 
@@ -120,7 +150,7 @@ async def get(self, archive_path, include_body=False):
         self.set_header("content-disposition", "attachment; filename={}".format(archive_filename))
 
         self.canceled = False
-        self.flush_cb = ioloop.PeriodicCallback(self.flush, ARCHIVE_DOWNLOAD_FLUSH_DELAY)
+        self.flush_cb = ioloop.PeriodicCallback(self.flush, self.archive_download_flush_delay)
         self.flush_cb.start()
 
         args = (