From 04391e0bc6136de490db3fb05fd4c9609163228f 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] extension stop hooks feedback

---
 jupyter_server/pytest_plugin.py            |  4 +--
 jupyter_server/serverapp.py                | 38 ++++++++++++++--------
 jupyter_server/tests/extension/test_app.py |  3 +-
 3 files changed, 28 insertions(+), 17 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 8c260a3f12..5f9718cb6f 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())
@@ -2084,7 +2084,7 @@ def cleanup_terminals(self):
         self.log.info(terminal_msg % n_terminals)
         run_sync(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(
@@ -2093,7 +2093,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"
@@ -2331,15 +2331,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()
+        await self.cleanup_extensions()
         self.cleanup_kernels()
         self.cleanup_terminals()
-        self.cleanup_extensions()
 
     def start_ioloop(self):
         """Start the IO Loop."""
@@ -2362,13 +2362,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