From de892c0d2880a2f3441af6fe0538282989758158 Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Fri, 21 May 2021 13:28:39 +0100
Subject: [PATCH] stop hook for extensions

* closes https://github.com/jupyter-server/jupyter_server/issues/241
* call a stop_extension method on server shutdown if present
---
 docs/source/developers/extensions.rst   |  5 +++++
 jupyter_server/extension/application.py |  2 ++
 jupyter_server/extension/manager.py     | 30 ++++++++++++++++++++++---
 jupyter_server/serverapp.py             |  5 +++++
 4 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/docs/source/developers/extensions.rst b/docs/source/developers/extensions.rst
index 878f0287e0..0fc19be459 100644
--- a/docs/source/developers/extensions.rst
+++ b/docs/source/developers/extensions.rst
@@ -156,6 +156,10 @@ The basic structure of an ExtensionApp is shown below:
             ...
             # Change the jinja templating environment
 
+        def stop_extension(self):
+            ...
+            # Perform any required shut down steps
+
 
 The ``ExtensionApp`` uses the following methods and properties to connect your extension to the Jupyter server. You do not need to define a ``_load_jupyter_server_extension`` function for these apps. Instead, overwrite the pieces below to add your custom settings, handlers and templates:
 
@@ -164,6 +168,7 @@ Methods
 * ``initialize_setting()``: adds custom settings to the Tornado Web Application.
 * ``initialize_handlers()``: appends handlers to the Tornado Web Application.
 * ``initialize_templates()``: initialize the templating engine (e.g. jinja2) for your frontend.
+* ``stop_extension()``: called on server shut down.
 
 Properties
 
diff --git a/jupyter_server/extension/application.py b/jupyter_server/extension/application.py
index e5cb6bcafd..abbabf17ac 100644
--- a/jupyter_server/extension/application.py
+++ b/jupyter_server/extension/application.py
@@ -419,6 +419,8 @@ def start(self):
     def stop(self):
         """Stop the underlying Jupyter server.
         """
+        if hasattr(self, 'stop_extension'):
+            self.stop_extension()
         self.serverapp.stop()
         self.serverapp.clear_instance()
 
diff --git a/jupyter_server/extension/manager.py b/jupyter_server/extension/manager.py
index d3f40d4da6..3c93f1950f 100644
--- a/jupyter_server/extension/manager.py
+++ b/jupyter_server/extension/manager.py
@@ -230,15 +230,17 @@ def link_point(self, point_name, serverapp):
 
     def load_point(self, point_name, serverapp):
         point = self.extension_points[point_name]
-        point.load(serverapp)
+        return point.load(serverapp)
 
     def link_all_points(self, serverapp):
         for point_name in self.extension_points:
             self.link_point(point_name, serverapp)
 
     def load_all_points(self, serverapp):
-        for point_name in self.extension_points:
+        return [
             self.load_point(point_name, serverapp)
+            for point_name in self.extension_points
+        ]
 
 
 class ExtensionManager(LoggingConfigurable):
@@ -285,6 +287,13 @@ def _config_manager_changed(self, change):
         """
     )
 
+    extension_apps = Dict(
+        help="""
+        Dictionary with extension names as keys
+        and ExtensionApp objects as values.
+        """
+    )
+
     @property
     def extension_points(self):
         extensions = self.extensions
@@ -338,12 +347,21 @@ def load_extension(self, name, serverapp):
         extension = self.extensions.get(name)
         if extension.enabled:
             try:
-                extension.load_all_points(serverapp)
+                self.extension_apps.setdefault(name, []).extend(
+                    extension.load_all_points(serverapp)
+                )
                 self.log.info("{name} | extension was successfully loaded.".format(name=name))
             except Exception as e:
                 self.log.debug("".join(traceback.format_exception(*sys.exc_info())))
                 self.log.warning("{name} | extension failed loading with message: {error}".format(name=name,error=str(e)))
 
+    def stop_extension(self, name, apps):
+        """Call the shutdown hooks in the specified apps."""
+        self.log.debug('stopping extension: {name}'.format(name=name))
+        for app in apps:
+            if hasattr(app, 'stop_extension'):
+                app.stop_extension()
+
     def link_all_extensions(self, serverapp):
         """Link all enabled extensions
         to an instance of ServerApp
@@ -361,3 +379,9 @@ def load_all_extensions(self, serverapp):
         # order.
         for name in sorted(self.extensions.keys()):
             self.load_extension(name, serverapp)
+
+    def stop_all_extensions(self, serverapp):
+        """Call the shutdown hooks in all extensions."""
+        for name, apps in sorted(dict(self.extension_apps).items()):
+            self.stop_extension(name, apps)
+            del self.extension_apps[name]
diff --git a/jupyter_server/serverapp.py b/jupyter_server/serverapp.py
index e3e404e541..d5846ff761 100755
--- a/jupyter_server/serverapp.py
+++ b/jupyter_server/serverapp.py
@@ -1921,6 +1921,10 @@ def cleanup_terminals(self):
         self.log.info(terminal_msg % n_terminals)
         run_sync(terminal_manager.terminate_all())
 
+    def cleanup_extensions(self):
+        """Call shutdown hooks in all extensions."""
+        self.extension_manager.stop_all_extensions(self)
+
     def running_server_info(self, kernel_count=True):
         "Return the current working directory and the server url information"
         info = self.contents_manager.info_string() + "\n"
@@ -2152,6 +2156,7 @@ def _cleanup(self):
         self.remove_browser_open_files()
         self.cleanup_kernels()
         self.cleanup_terminals()
+        self.cleanup_extensions()
 
     def start_ioloop(self):
         """Start the IO Loop."""