From ed2b3e06b6dcb3520072851e2eaae9bcc9764120 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 1/7] 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     | 29 ++++++++++++++++++++++---
 jupyter_server/serverapp.py             | 12 ++++++++++
 4 files changed, 45 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 3e0d8e1bd6..c7aa8bc6c4 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):
@@ -290,6 +292,13 @@ def sorted_extensions(self):
         """
     )
 
+    extension_apps = Dict(
+        help="""
+        Dictionary with extension names as keys
+        and ExtensionApp objects as values.
+        """
+    )
+
     @property
     def extension_points(self):
         extensions = self.extensions
@@ -343,12 +352,20 @@ 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."""
+        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
@@ -366,3 +383,9 @@ def load_all_extensions(self, serverapp):
         # order.
         for name in self.sorted_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 dd9939ca5d..e69b2338da 100755
--- a/jupyter_server/serverapp.py
+++ b/jupyter_server/serverapp.py
@@ -2085,6 +2085,17 @@ 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."""
+        n_extensions = len(self.extension_manager.extension_apps)
+        extension_msg = trans.ngettext(
+            'Shutting down %d extension',
+            'Shutting down %d extensions',
+            n_extensions
+        )
+        self.log.info(extension_msg % n_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"
@@ -2329,6 +2340,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."""

From 0c5f8ca9de173e3f6bb5eb4f9903ae8a4748aa72 Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Fri, 11 Jun 2021 12:53:30 +0100
Subject: [PATCH 2/7] fix typo

---
 docs/source/developers/extensions.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/developers/extensions.rst b/docs/source/developers/extensions.rst
index 0fc19be459..73b918d4ce 100644
--- a/docs/source/developers/extensions.rst
+++ b/docs/source/developers/extensions.rst
@@ -165,7 +165,7 @@ The ``ExtensionApp`` uses the following methods and properties to connect your e
 
 Methods
 
-* ``initialize_setting()``: adds custom settings to the Tornado Web Application.
+* ``initialize_settings()``: 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.

From 584773d57b3ba463e18fe5f6b2b028e90db11c5a Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Fri, 11 Jun 2021 15:37:39 +0100
Subject: [PATCH 3/7] make extension stop hooks async

---
 docs/source/developers/extensions.rst |  2 +-
 jupyter_server/extension/manager.py   | 16 +++++++++++-----
 jupyter_server/serverapp.py           |  4 +---
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/docs/source/developers/extensions.rst b/docs/source/developers/extensions.rst
index 73b918d4ce..b056746443 100644
--- a/docs/source/developers/extensions.rst
+++ b/docs/source/developers/extensions.rst
@@ -156,7 +156,7 @@ The basic structure of an ExtensionApp is shown below:
             ...
             # Change the jinja templating environment
 
-        def stop_extension(self):
+        async def stop_extension(self):
             ...
             # Perform any required shut down steps
 
diff --git a/jupyter_server/extension/manager.py b/jupyter_server/extension/manager.py
index c7aa8bc6c4..42c24196da 100644
--- a/jupyter_server/extension/manager.py
+++ b/jupyter_server/extension/manager.py
@@ -2,6 +2,8 @@
 import sys
 import traceback
 
+from tornado.gen import multi
+
 from traitlets.config import LoggingConfigurable
 
 from traitlets import (
@@ -360,11 +362,14 @@ def load_extension(self, name, serverapp):
                 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):
+    async def stop_extension(self, name, apps):
         """Call the shutdown hooks in the specified apps."""
         for app in apps:
             if hasattr(app, 'stop_extension'):
-                app.stop_extension()
+                await app.stop_extension()
+            if name in self.extension_apps:
+                # might not be the case in tests
+                del self.extension_apps[name]
 
     def link_all_extensions(self, serverapp):
         """Link all enabled extensions
@@ -384,8 +389,9 @@ def load_all_extensions(self, serverapp):
         for name in self.sorted_extensions.keys():
             self.load_extension(name, serverapp)
 
-    def stop_all_extensions(self, serverapp):
+    async def stop_all_extensions(self, serverapp):
         """Call the shutdown hooks in all extensions."""
-        for name, apps in sorted(dict(self.extension_apps).items()):
+        await multi([
             self.stop_extension(name, apps)
-            del self.extension_apps[name]
+            for name, apps in sorted(dict(self.extension_apps).items())
+        ])
diff --git a/jupyter_server/serverapp.py b/jupyter_server/serverapp.py
index e69b2338da..a18247cc04 100755
--- a/jupyter_server/serverapp.py
+++ b/jupyter_server/serverapp.py
@@ -2094,7 +2094,7 @@ def cleanup_extensions(self):
             n_extensions
         )
         self.log.info(extension_msg % n_extensions)
-        self.extension_manager.stop_all_extensions(self)
+        run_sync(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"
@@ -2353,8 +2353,6 @@ def start_ioloop(self):
             self.io_loop.start()
         except KeyboardInterrupt:
             self.log.info(_i18n("Interrupted..."))
-        finally:
-            self._cleanup()
 
     def init_ioloop(self):
         """init self.io_loop so that an extension can use it by io_loop.call_later() to create background tasks"""

From 4801d823c7a8916d20689492fc02dbd61135c483 Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Fri, 11 Jun 2021 16:40:22 +0100
Subject: [PATCH 4/7] extension stop hook tests

---
 jupyter_server/extension/manager.py        | 16 ++++++---
 jupyter_server/tests/extension/test_app.py | 39 ++++++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/jupyter_server/extension/manager.py b/jupyter_server/extension/manager.py
index 42c24196da..62d9c95c48 100644
--- a/jupyter_server/extension/manager.py
+++ b/jupyter_server/extension/manager.py
@@ -297,7 +297,7 @@ def sorted_extensions(self):
     extension_apps = Dict(
         help="""
         Dictionary with extension names as keys
-        and ExtensionApp objects as values.
+        and sets of ExtensionApp objects as values.
         """
     )
 
@@ -352,15 +352,21 @@ def link_extension(self, name, serverapp):
 
     def load_extension(self, name, serverapp):
         extension = self.extensions.get(name)
+
         if extension.enabled:
             try:
-                self.extension_apps.setdefault(name, []).extend(
-                    extension.load_all_points(serverapp)
-                )
-                self.log.info("{name} | extension was successfully loaded.".format(name=name))
+                points = extension.load_all_points(serverapp)
             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)))
+            else:
+                self.extension_apps.setdefault(name, set()).update((
+                    point
+                    for point in points
+                    if point is not None
+                ))
+                self.log.info("{name} | extension was successfully loaded.".format(name=name))
+
 
     async def stop_extension(self, name, apps):
         """Call the shutdown hooks in the specified apps."""
diff --git a/jupyter_server/tests/extension/test_app.py b/jupyter_server/tests/extension/test_app.py
index 3cc0e82fe6..25e54cfcfc 100644
--- a/jupyter_server/tests/extension/test_app.py
+++ b/jupyter_server/tests/extension/test_app.py
@@ -101,3 +101,42 @@ def test_load_parallel_extensions(monkeypatch, jp_environ):
     exts = serverapp.jpserver_extensions
     assert exts['jupyter_server.tests.extension.mockextensions.mock1']
     assert exts['jupyter_server.tests.extension.mockextensions']
+
+
+def test_stop_extension(jp_serverapp, caplog):
+    """Test the stop_extension method.
+
+    This should be fired by ServerApp.cleanup_extensions.
+    """
+    calls = 0
+
+    # load extensions (make sure we only have the one extension loaded
+    jp_serverapp.extension_manager.load_all_extensions(jp_serverapp)
+    assert list(jp_serverapp.extension_manager.extension_apps) == [
+        'jupyter_server.tests.extension.mockextensions'
+    ]
+
+    # add a stop_extension method for the extension app
+    async def _stop(*args):
+        nonlocal calls
+        calls += 1
+    for apps in jp_serverapp.extension_manager.extension_apps.values():
+        for app in apps:
+            if app:
+                app.stop_extension = _stop
+
+    # call cleanup_extensions, check the logging is correct
+    caplog.clear()
+    jp_serverapp.cleanup_extensions()
+    assert [
+        msg
+        for *_, msg in caplog.record_tuples
+    ] == [
+        'Shutting down 1 extension'
+    ]
+
+    # check the extension_apps dictionary is updated
+    assert list(jp_serverapp.extension_manager.extension_apps) == []
+
+    # check the shutdown method was called once
+    assert calls == 1

From 357b9e5aee5ca32580c3a3c18c03f9b42a78fe55 Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Tue, 29 Jun 2021 16:18:58 +0100
Subject: [PATCH 5/7] extension stop hooks feedback

---
 jupyter_server/pytest_plugin.py            |  4 +-
 jupyter_server/serverapp.py                | 50 +++++++++++++---------
 jupyter_server/tests/extension/test_app.py |  3 +-
 3 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/jupyter_server/pytest_plugin.py b/jupyter_server/pytest_plugin.py
index 892828ba5f..00a7c8fc86 100644
--- a/jupyter_server/pytest_plugin.py
+++ b/jupyter_server/pytest_plugin.py
@@ -18,7 +18,7 @@
 
 from jupyter_server.extension import serverextension
 from jupyter_server.serverapp import ServerApp
-from jupyter_server.utils import url_path_join
+from jupyter_server.utils import url_path_join, run_sync
 from jupyter_server.services.contents.filemanager import FileContentsManager
 from jupyter_server.services.contents.largefilemanager import LargeFileManager
 
@@ -284,7 +284,7 @@ def jp_serverapp(
     """Starts a Jupyter Server instance based on the established configuration values."""
     app = jp_configurable_serverapp(config=jp_server_config, argv=jp_argv)
     yield app
-    app._cleanup()
+    run_sync(app._cleanup())
 
 
 @pytest.fixture
diff --git a/jupyter_server/serverapp.py b/jupyter_server/serverapp.py
index a18247cc04..0e4f4f41c9 100755
--- a/jupyter_server/serverapp.py
+++ b/jupyter_server/serverapp.py
@@ -1750,7 +1750,7 @@ def _confirm_exit(self):
             self.log.critical(_i18n("Shutting down..."))
             # schedule stop on the main thread,
             # since this might be called from a signal handler
-            self.io_loop.add_callback_from_signal(self.io_loop.stop)
+            self.stop(from_signal=True)
             return
         print(self.running_server_info())
         yes = _i18n('y')
@@ -1764,7 +1764,7 @@ def _confirm_exit(self):
                 self.log.critical(_i18n("Shutdown confirmed"))
                 # schedule stop on the main thread,
                 # since this might be called from a signal handler
-                self.io_loop.add_callback_from_signal(self.io_loop.stop)
+                self.stop(from_signal=True)
                 return
         else:
             print(_i18n("No answer for 5s:"), end=' ')
@@ -1777,7 +1777,7 @@ def _confirm_exit(self):
 
     def _signal_stop(self, sig, frame):
         self.log.critical(_i18n("received signal %s, stopping"), sig)
-        self.io_loop.add_callback_from_signal(self.io_loop.stop)
+        self.stop(from_signal=True)
 
     def _signal_info(self, sig, frame):
         print(self.running_server_info())
@@ -2059,7 +2059,7 @@ def initialize(self, argv=None, find_extensions=True, new_httpserver=True, start
         if new_httpserver:
             self.init_httpserver()
 
-    def cleanup_kernels(self):
+    async def cleanup_kernels(self):
         """Shutdown all kernels.
 
         The kernels will shutdown themselves when this process no longer exists,
@@ -2068,9 +2068,9 @@ def cleanup_kernels(self):
         n_kernels = len(self.kernel_manager.list_kernel_ids())
         kernel_msg = trans.ngettext('Shutting down %d kernel', 'Shutting down %d kernels', n_kernels)
         self.log.info(kernel_msg % n_kernels)
-        run_sync(self.kernel_manager.shutdown_all())
+        await self.kernel_manager.shutdown_all()
 
-    def cleanup_terminals(self):
+    async def cleanup_terminals(self):
         """Shutdown all terminals.
 
         The terminals will shutdown themselves when this process no longer exists,
@@ -2083,9 +2083,9 @@ def cleanup_terminals(self):
         n_terminals = len(terminal_manager.list())
         terminal_msg = trans.ngettext('Shutting down %d terminal', 'Shutting down %d terminals', n_terminals)
         self.log.info(terminal_msg % n_terminals)
-        run_sync(terminal_manager.terminate_all())
+        await terminal_manager.terminate_all()
 
-    def cleanup_extensions(self):
+    async def cleanup_extensions(self):
         """Call shutdown hooks in all extensions."""
         n_extensions = len(self.extension_manager.extension_apps)
         extension_msg = trans.ngettext(
@@ -2094,7 +2094,7 @@ def cleanup_extensions(self):
             n_extensions
         )
         self.log.info(extension_msg % n_extensions)
-        run_sync(self.extension_manager.stop_all_extensions(self))
+        await 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"
@@ -2332,15 +2332,15 @@ def start_app(self):
                     '    %s' % self.display_url,
                 ]))
 
-    def _cleanup(self):
-        """General cleanup of files and kernels created
+    async def _cleanup(self):
+        """General cleanup of files, extensions and kernels created
         by this instance ServerApp.
         """
         self.remove_server_info_file()
         self.remove_browser_open_files()
-        self.cleanup_kernels()
-        self.cleanup_terminals()
-        self.cleanup_extensions()
+        await self.cleanup_extensions()
+        await self.cleanup_kernels()
+        await self.cleanup_terminals()
 
     def start_ioloop(self):
         """Start the IO Loop."""
@@ -2366,13 +2366,23 @@ def start(self):
         self.start_app()
         self.start_ioloop()
 
-    def stop(self):
-        def _stop():
+    async def _stop(self):
+        """Cleanup resources and stop the IO Loop."""
+        await self._cleanup()
+        self.io_loop.stop()
+
+    def stop(self, from_signal=False):
+        """Cleanup resources and stop the server."""
+        if hasattr(self, '_http_server'):
             # Stop a server if its set.
-            if hasattr(self, '_http_server'):
-                self.http_server.stop()
-            self.io_loop.stop()
-        self.io_loop.add_callback(_stop)
+            self.http_server.stop()
+        if getattr(self, 'io_loop', None):
+            # use IOLoop.add_callback because signal.signal must be called
+            # from main thread
+            if from_signal:
+                self.io_loop.add_callback_from_signal(self._stop)
+            else:
+                self.io_loop.add_callback(self._stop)
 
 
 def list_running_servers(runtime_dir=None):
diff --git a/jupyter_server/tests/extension/test_app.py b/jupyter_server/tests/extension/test_app.py
index 25e54cfcfc..4ee4fb11ed 100644
--- a/jupyter_server/tests/extension/test_app.py
+++ b/jupyter_server/tests/extension/test_app.py
@@ -1,6 +1,7 @@
 import pytest
 from traitlets.config import Config
 from jupyter_server.serverapp import ServerApp
+from jupyter_server.utils import run_sync
 from .mockextensions.app import MockExtensionApp
 
 
@@ -127,7 +128,7 @@ async def _stop(*args):
 
     # call cleanup_extensions, check the logging is correct
     caplog.clear()
-    jp_serverapp.cleanup_extensions()
+    run_sync(jp_serverapp.cleanup_extensions())
     assert [
         msg
         for *_, msg in caplog.record_tuples

From 42e3f02cd87f38df36292aa1916d255f5a1caf13 Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Thu, 1 Jul 2021 10:42:05 +0100
Subject: [PATCH 6/7] run_sync

---
 jupyter_server/serverapp.py | 10 ++++++----
 jupyter_server/utils.py     | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/jupyter_server/serverapp.py b/jupyter_server/serverapp.py
index 0e4f4f41c9..03fd748902 100755
--- a/jupyter_server/serverapp.py
+++ b/jupyter_server/serverapp.py
@@ -43,7 +43,7 @@
 
 from jupyter_core.paths import secure_write
 from jupyter_server.transutils import trans, _i18n
-from jupyter_server.utils import run_sync
+from jupyter_server.utils import run_sync_in_loop
 
 # the minimum viable tornado version: needs to be kept in sync with setup.py
 MIN_TORNADO = (6, 1, 0)
@@ -2068,7 +2068,7 @@ async def cleanup_kernels(self):
         n_kernels = len(self.kernel_manager.list_kernel_ids())
         kernel_msg = trans.ngettext('Shutting down %d kernel', 'Shutting down %d kernels', n_kernels)
         self.log.info(kernel_msg % n_kernels)
-        await self.kernel_manager.shutdown_all()
+        await run_sync_in_loop(self.kernel_manager.shutdown_all())
 
     async def cleanup_terminals(self):
         """Shutdown all terminals.
@@ -2083,7 +2083,7 @@ async def cleanup_terminals(self):
         n_terminals = len(terminal_manager.list())
         terminal_msg = trans.ngettext('Shutting down %d terminal', 'Shutting down %d terminals', n_terminals)
         self.log.info(terminal_msg % n_terminals)
-        await terminal_manager.terminate_all()
+        await run_sync_in_loop(terminal_manager.terminate_all())
 
     async def cleanup_extensions(self):
         """Call shutdown hooks in all extensions."""
@@ -2094,7 +2094,9 @@ async def cleanup_extensions(self):
             n_extensions
         )
         self.log.info(extension_msg % n_extensions)
-        await self.extension_manager.stop_all_extensions(self)
+        await run_sync_in_loop(
+            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"
diff --git a/jupyter_server/utils.py b/jupyter_server/utils.py
index e0a532bbdd..83fdd43611 100644
--- a/jupyter_server/utils.py
+++ b/jupyter_server/utils.py
@@ -230,6 +230,29 @@ def wrapped():
     return wrapped()
 
 
+async def run_sync_in_loop(maybe_async):
+    """Runs a function synchronously whether it is an async function or not.
+
+    If async, runs maybe_async and blocks until it has executed.
+
+    If not async, just returns maybe_async as it is the result of something
+    that has already executed.
+
+    Parameters
+    ----------
+    maybe_async : async or non-async object
+        The object to be executed, if it is async.
+
+    Returns
+    -------
+    result
+        Whatever the async object returns, or the object itself.
+    """
+    if not inspect.isawaitable(maybe_async):
+        return maybe_async
+    return await maybe_async
+
+
 def urlencode_unix_socket_path(socket_path):
     """Encodes a UNIX socket path string from a socket path for the `http+unix` URI form."""
     return socket_path.replace('/', '%2F')

From 0eeeff8e8bd18901c4020c938f13d5cb11ee7fab Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Wed, 7 Jul 2021 13:51:18 +0100
Subject: [PATCH 7/7] extension stop hooks extension_apps property

---
 jupyter_server/extension/application.py    |  5 ++--
 jupyter_server/extension/manager.py        | 35 +++++++++++-----------
 jupyter_server/tests/extension/test_app.py | 10 +++----
 3 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/jupyter_server/extension/application.py b/jupyter_server/extension/application.py
index abbabf17ac..76ef17dd07 100644
--- a/jupyter_server/extension/application.py
+++ b/jupyter_server/extension/application.py
@@ -416,11 +416,12 @@ def start(self):
         # Start the server.
         self.serverapp.start()
 
+    async def stop_extension(self):
+        """Cleanup any resources managed by this extension."""
+
     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 62d9c95c48..83f1af7943 100644
--- a/jupyter_server/extension/manager.py
+++ b/jupyter_server/extension/manager.py
@@ -294,19 +294,26 @@ def sorted_extensions(self):
         """
     )
 
-    extension_apps = Dict(
-        help="""
-        Dictionary with extension names as keys
-        and sets of ExtensionApp objects as values.
+    @property
+    def extension_apps(self):
+        """Return mapping of extension names and sets of ExtensionApp objects.
         """
-    )
+        return {
+            name: {
+                point.app
+                for point in extension.extension_points.values()
+                if point.app
+            }
+            for name, extension in self.extensions.items()
+        }
 
     @property
     def extension_points(self):
-        extensions = self.extensions
+        """Return mapping of extension point names and ExtensionPoint objects.
+        """
         return {
             name: point
-            for value in extensions.values()
+            for value in self.extensions.values()
             for name, point in value.extension_points.items()
         }
 
@@ -360,22 +367,14 @@ def load_extension(self, name, serverapp):
                 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)))
             else:
-                self.extension_apps.setdefault(name, set()).update((
-                    point
-                    for point in points
-                    if point is not None
-                ))
                 self.log.info("{name} | extension was successfully loaded.".format(name=name))
 
-
     async def stop_extension(self, name, apps):
         """Call the shutdown hooks in the specified apps."""
         for app in apps:
-            if hasattr(app, 'stop_extension'):
-                await app.stop_extension()
-            if name in self.extension_apps:
-                # might not be the case in tests
-                del self.extension_apps[name]
+            self.log.debug('{} | extension app "{}" stopping'.format(name, app.name))
+            await app.stop_extension()
+            self.log.debug('{} | extension app "{}" stopped'.format(name, app.name))
 
     def link_all_extensions(self, serverapp):
         """Link all enabled extensions
diff --git a/jupyter_server/tests/extension/test_app.py b/jupyter_server/tests/extension/test_app.py
index 4ee4fb11ed..fe83d24ba2 100644
--- a/jupyter_server/tests/extension/test_app.py
+++ b/jupyter_server/tests/extension/test_app.py
@@ -113,8 +113,9 @@ def test_stop_extension(jp_serverapp, caplog):
 
     # load extensions (make sure we only have the one extension loaded
     jp_serverapp.extension_manager.load_all_extensions(jp_serverapp)
+    extension_name = 'jupyter_server.tests.extension.mockextensions'
     assert list(jp_serverapp.extension_manager.extension_apps) == [
-        'jupyter_server.tests.extension.mockextensions'
+        extension_name
     ]
 
     # add a stop_extension method for the extension app
@@ -133,11 +134,10 @@ async def _stop(*args):
         msg
         for *_, msg in caplog.record_tuples
     ] == [
-        'Shutting down 1 extension'
+        'Shutting down 1 extension',
+        '{} | extension app "mockextension" stopping'.format(extension_name),
+        '{} | extension app "mockextension" stopped'.format(extension_name),
     ]
 
-    # check the extension_apps dictionary is updated
-    assert list(jp_serverapp.extension_manager.extension_apps) == []
-
     # check the shutdown method was called once
     assert calls == 1