Add async start hook to ExtensionApp API

Author: Zsailer

jupyter_server/extension/application.py

@@ -446,6 +446,9 @@ def start(self):
         assert self.serverapp is not None
         self.serverapp.start()
 
+    async def start_extension(self):
+        """An async hook to start e.g. tasks after the server's event loop is running."""
+
     def current_activity(self):
         """Return a list of activity happening in this extension."""
         return

jupyter_server/extension/manager.py

@@ -367,6 +367,16 @@ def load_extension(self, name):
             else:
                 self.log.info("%s | extension was successfully loaded.", name)
 
+    async def start_extension(self, name, apps):
+        """Call the start hooks in the specified apps."""
+        for app in apps:
+            self.log.debug("%s | extension app %r starting", name, app.name)
+            try:
+                await app.start_extension()
+                self.log.debug("%s | extension app %r started", name, app.name)
+            except Exception as err:
+                self.log.error(err)
+
     async def stop_extension(self, name, apps):
         """Call the shutdown hooks in the specified apps."""
         for app in apps:
@@ -392,6 +402,10 @@ def load_all_extensions(self):
         for name in self.sorted_extensions:
             self.load_extension(name)
 
+    async def start_all_extensions(self):
+        """Call the start hooks in all extensions."""
+        await multi(list(starmap(self.start_extension, sorted(dict(self.extension_apps).items()))))
+
     async def stop_all_extensions(self):
         """Call the shutdown hooks in all extensions."""
         await multi(list(starmap(self.stop_extension, sorted(dict(self.extension_apps).items()))))

jupyter_server/serverapp.py

@@ -3060,6 +3060,72 @@ def start_app(self) -> None:
                 )
                 self.exit(1)
 
+        self.write_server_info_file()
+
+        if not self.no_browser_open_file:
+            self.write_browser_open_files()
+
+        # Handle the browser opening.
+        if self.open_browser and not self.sock:
+            self.launch_browser()
+
+    async def _cleanup(self) -> None:
+        """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()
+        await self.cleanup_kernels()
+        try:
+            await self.kernel_websocket_connection_class.close_all()  # type:ignore[attr-defined]
+        except AttributeError:
+            # This can happen in two different scenarios:
+            #
+            # 1. During tests, where the _cleanup method is invoked without
+            #    the corresponding initialize method having been invoked.
+            # 2. If the provided `kernel_websocket_connection_class` does not
+            #    implement the `close_all` class method.
+            #
+            # In either case, we don't need to do anything and just want to treat
+            # the raised error as a no-op.
+            pass
+        if getattr(self, "kernel_manager", None):
+            self.kernel_manager.__del__()
+        if getattr(self, "session_manager", None):
+            self.session_manager.close()
+        if hasattr(self, "http_server"):
+            # Stop a server if its set.
+            self.http_server.stop()
+
+    def start_ioloop(self) -> None:
+        """Start the IO Loop."""
+        if sys.platform.startswith("win"):
+            # add no-op to wake every 5s
+            # to handle signals that may be ignored by the inner loop
+            pc = ioloop.PeriodicCallback(lambda: None, 5000)
+            pc.start()
+        try:
+            self.io_loop.add_callback(self._post_start)
+            self.io_loop.start()
+        except KeyboardInterrupt:
+            self.log.info(_i18n("Interrupted..."))
+
+    def init_ioloop(self) -> None:
+        """init self.io_loop so that an extension can use it by io_loop.call_later() to create background tasks"""
+        self.io_loop = ioloop.IOLoop.current()
+
+    async def _post_start(self):
+        """Add an async hook to start tasks after the event loop is running.
+
+        This will also attempt to start all tasks found in
+        the `start_extension` method in Extension Apps.
+        """
+        try:
+            await self.extension_manager.start_all_extensions()
+        except Exception as err:
+            self.log.error(err)
+
         info = self.log.info
         for line in self.running_server_info(kernel_count=False).split("\n"):
             info(line)
@@ -3078,15 +3144,6 @@ def start_app(self) -> None:
                 )
             )
 
-        self.write_server_info_file()
-
-        if not self.no_browser_open_file:
-            self.write_browser_open_files()
-
-        # Handle the browser opening.
-        if self.open_browser and not self.sock:
-            self.launch_browser()
-
         if self.identity_provider.token and self.identity_provider.token_generated:
             # log full URL with generated token, so there's a copy/pasteable link
             # with auth info.
@@ -3127,51 +3184,6 @@ def start_app(self) -> None:
 
                 self.log.critical("\n".join(message))
 
-    async def _cleanup(self) -> None:
-        """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()
-        await self.cleanup_kernels()
-        try:
-            await self.kernel_websocket_connection_class.close_all()  # type:ignore[attr-defined]
-        except AttributeError:
-            # This can happen in two different scenarios:
-            #
-            # 1. During tests, where the _cleanup method is invoked without
-            #    the corresponding initialize method having been invoked.
-            # 2. If the provided `kernel_websocket_connection_class` does not
-            #    implement the `close_all` class method.
-            #
-            # In either case, we don't need to do anything and just want to treat
-            # the raised error as a no-op.
-            pass
-        if getattr(self, "kernel_manager", None):
-            self.kernel_manager.__del__()
-        if getattr(self, "session_manager", None):
-            self.session_manager.close()
-        if hasattr(self, "http_server"):
-            # Stop a server if its set.
-            self.http_server.stop()
-
-    def start_ioloop(self) -> None:
-        """Start the IO Loop."""
-        if sys.platform.startswith("win"):
-            # add no-op to wake every 5s
-            # to handle signals that may be ignored by the inner loop
-            pc = ioloop.PeriodicCallback(lambda: None, 5000)
-            pc.start()
-        try:
-            self.io_loop.start()
-        except KeyboardInterrupt:
-            self.log.info(_i18n("Interrupted..."))
-
-    def init_ioloop(self) -> None:
-        """init self.io_loop so that an extension can use it by io_loop.call_later() to create background tasks"""
-        self.io_loop = ioloop.IOLoop.current()
-
     def start(self) -> None:
         """Start the Jupyter server app, after initialization
 

tests/extension/mockextensions/app.py

@@ -5,7 +5,7 @@
 from jupyter_events import EventLogger
 from jupyter_events.schema_registry import SchemaRegistryException
 from tornado import web
-from traitlets import List, Unicode
+from traitlets import Bool, List, Unicode
 
 from jupyter_server.base.handlers import JupyterHandler
 from jupyter_server.extension.application import ExtensionApp, ExtensionAppJinjaMixin
@@ -56,6 +56,7 @@ class MockExtensionApp(ExtensionAppJinjaMixin, ExtensionApp):
     static_paths = [STATIC_PATH]  # type:ignore[assignment]
     mock_trait = Unicode("mock trait", config=True)
     loaded = False
+    started = Bool(False)
 
     serverapp_config = {
         "jpserver_extensions": {
@@ -96,6 +97,9 @@ def initialize_handlers(self):
         self.handlers.append(("/mock_error_notemplate", MockExtensionErrorHandler))
         self.loaded = True
 
+    async def start_extension(self):
+        self.started = True
+
 
 if __name__ == "__main__":
     MockExtensionApp.launch_instance()

tests/extension/test_app.py

@@ -140,6 +140,11 @@ async def test_load_parallel_extensions(monkeypatch, jp_environ):
     assert exts["tests.extension.mockextensions"]
 
 
+async def test_start_extension(jp_serverapp, mock_extension):
+    await jp_serverapp._post_start()
+    assert mock_extension.started
+
+
 async def test_stop_extension(jp_serverapp, caplog):
     """Test the stop_extension method.
 

tests/test_serverapp.py

@@ -607,8 +607,9 @@ def test_running_server_info(jp_serverapp):
 
 
 @pytest.mark.parametrize("should_exist", [True, False])
-def test_browser_open_files(jp_configurable_serverapp, should_exist, caplog):
+async def test_browser_open_files(jp_configurable_serverapp, should_exist, caplog):
     app = jp_configurable_serverapp(no_browser_open_file=not should_exist)
+    await app._post_start()
     assert os.path.exists(app.browser_open_file) == should_exist
     url = urljoin("file:", pathname2url(app.browser_open_file))
     url_messages = [rec.message for rec in caplog.records if url in rec.message]