Merged PR posit-dev/positron-python#109: improve kernel and language server shutdown behaviour

Merge pull request #109 from posit-dev/improve-shutdown-behaviour

improve kernel and language server shutdown behaviour
--------------------
Commit message for posit-dev/positron-python@7750de1:

black formatting

--------------------
Commit message for posit-dev/positron-python@1efb00d:

improve kernel and language server shutdown behaviour


Authored-by: Wasim Lorgat <[email protected]>
Signed-off-by: Wasim Lorgat <[email protected]>

Author: seeM

extensions/positron-python/pythonFiles/positron/positron_ipkernel.py

@@ -31,12 +31,18 @@
 POSITRON_PLOT_COMM = "positron.plot"
 """The comm channel target_name for Positron's Plots View"""
 
-POSITON_NS_HIDDEN = {"_exit_code": {}, "__pydevd_ret_val_dict": {}, "__warningregistry__": {}}
+POSITON_NS_HIDDEN = {
+    "_exit_code": {},
+    "__pydevd_ret_val_dict": {},
+    "__warningregistry__": {},
+}
 """Additional variables to hide from the user's namespace."""
 
 # Key used to store the user's environment snapshot in the hidden namespace
 __POSITRON_CACHE_KEY__ = "__positron_cache__"
 
+logger = logging.getLogger(__name__)
+
 
 class PositronIPyKernel(IPythonKernel):
     """
@@ -75,12 +81,16 @@ def do_shutdown(self, restart) -> dict:
         """
         Handle kernel shutdown.
         """
-        result = super().do_shutdown(restart)
+        logger.info("Shutting down the kernel")
         self.display_pub_hook.shutdown()
         self.env_service.shutdown()
         self.lsp_service.shutdown()
         self.dataviewer_service.shutdown()
-        return result
+
+        # We don't call super().do_shutdown since it sets shell.exit_now = True which tries to
+        # stop the event loop at the same time as self.shutdown_request (since self.shell_stream.io_loop
+        # points to the same underlying asyncio loop).
+        return dict(status="ok", restart=restart)
 
     def handle_pre_execute(self) -> None:
         """

extensions/positron-python/pythonFiles/positron/positron_jedilsp.py

@@ -5,6 +5,8 @@
 import os
 import sys
 
+from pygls.protocol import lsp_method
+
 # Add the lib path to our sys path so jedi_language_server can find its references
 EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, os.path.join(EXTENSION_ROOT, "pythonFiles", "lib", "jedilsp"))
@@ -38,6 +40,7 @@
 )
 from lsprotocol.types import (
     COMPLETION_ITEM_RESOLVE,
+    EXIT,
     TEXT_DOCUMENT_CODE_ACTION,
     TEXT_DOCUMENT_COMPLETION,
     TEXT_DOCUMENT_DEFINITION,
@@ -87,6 +90,34 @@
     from .positron_ipkernel import PositronIPyKernel
 
 
+logger = logging.getLogger(__name__)
+
+
+class PositronJediLanguageServerProtocol(JediLanguageServerProtocol):
+    """Positron extension to the Jedi language server protocol."""
+
+    @lsp_method(EXIT)
+    def lsp_exit(self, *args) -> None:
+        """Override superclass to avoid system exit."""
+        if self.transport is not None:
+            self.transport.close()
+
+        # --- Start Positron
+        # Instead of calling sys.exit (which pygls catches to call self.shutdown) call shutdown directly.
+        asyncio.create_task(POSITRON.shutdown())
+        # --- End Positron
+
+    def connection_lost(self, exc):
+        """Override superclass to avoid system exit."""
+        logger.info("Connection lost")
+        if exc is not None:
+            logger.error("Connection lost with error:", exc_info=exc)
+
+    def eof_received(self):
+        """Override superclass to avoid closing on eof (e.g. when the Positron browser is refreshed)."""
+        return True
+
+
 class PositronJediLanguageServer(JediLanguageServer):
     """Positron extension to the Jedi language server."""
 
@@ -121,11 +152,7 @@ def start(self, lsp_host: str, lsp_port: int, kernel: "PositronIPyKernel") -> No
         completions with awareness of live variables from user's namespace.
         """
         self.kernel = kernel
-
-        try:
-            asyncio.ensure_future(self._start_jedi(lsp_host, lsp_port))
-        except KeyboardInterrupt:
-            pass
+        asyncio.create_task(self._start_jedi(lsp_host, lsp_port))
 
     async def _start_jedi(self, lsp_host, lsp_port):
         """
@@ -134,15 +161,36 @@ async def _start_jedi(self, lsp_host, lsp_port):
         Adapted from `pygls.server.Server.start_tcp` to use existing asyncio loop.
         """
         self._stop_event = Event()
+        self.lsp: PositronJediLanguageServerProtocol
         loop = asyncio.get_event_loop()
-        self._server = await loop.create_server(self.lsp, lsp_host, lsp_port)  # type: ignore
-        await self._server.serve_forever()
+        self._server = await loop.create_server(self.lsp, lsp_host, lsp_port)
+
+    async def shutdown(self):
+        """Override superclass to allow running async and avoid stopping the event loop."""
+        logger.info("Shutting down the language server")
+
+        self._stop_event.set()
+
+        if self._thread_pool:
+            self._thread_pool.terminate()
+            self._thread_pool.join()
+
+        if self._thread_pool_executor:
+            self._thread_pool_executor.shutdown()
+
+        if self._server:
+            self._server.close()
+            # --- Start Positron
+            await self._server.wait_closed()
+            # --- End Positron
+
+        logger.info("Language server is shut down")
 
 
 POSITRON = PositronJediLanguageServer(
     name="jedi-language-server",
     version="0.18.2",
-    protocol_cls=JediLanguageServerProtocol,
+    protocol_cls=PositronJediLanguageServerProtocol,
 )
 
 # Server Features
@@ -177,7 +225,7 @@ def positron_completion(
         ns = server.kernel.get_user_ns()
         namespaces.append(ns)
 
-    # Use Interpreter() to include the kernel namespaces in completions
+    # Use Interpreter instead of Script to include the kernel namespaces in completions
     jedi_script = Interpreter(
         document.source, namespaces, path=document.path, project=server.project
     )

extensions/positron-python/pythonFiles/positron_language_server.py

@@ -8,7 +8,6 @@
 import logging
 import os
 import sys
-import traceback
 
 from ipykernel import kernelapp
 
@@ -19,6 +18,8 @@
 
 from positron.positron_ipkernel import PositronIPyKernel
 
+logger = logging.getLogger(__name__)
+
 
 def initialize_config() -> None:
     """
@@ -85,39 +86,34 @@ def initialize_config() -> None:
             logging.warning(f"Unable to start debugpy: {error}", exc_info=True)
 
 
-async def start_ipykernel() -> None:
-    """Starts Positron's IPyKernel as the interpreter for our console."""
+if __name__ == "__main__":
+    exit_status = 0
+
+    # Init the configuration args
+    initialize_config()
+
+    # Start Positron's IPyKernel as the interpreter for our console.
     app = kernelapp.IPKernelApp.instance(kernel_class=PositronIPyKernel)
     app.initialize()
     app.kernel.start()
 
+    logger.info(f"Process ID {os.getpid()}")
 
-if __name__ == "__main__":
-    exitStatus = 0
+    # IPyKernel uses Tornado which (as of version 5.0) shares the same event
+    # loop as asyncio.
+    loop = asyncio.get_event_loop()
 
-    try:
-        # Init the configuration args
-        initialize_config()
+    # Enable asyncio debug mode.
+    if logging.getLogger().level == logging.DEBUG:
+        loop.set_debug(True)
 
-        # Start Positron's IPyKernel as the interpreter for our console.
-        loop = asyncio.get_event_loop()
-        try:
-            asyncio.ensure_future(start_ipykernel())
-            loop.run_forever()
-        except KeyboardInterrupt:
-            pass
-        finally:
-            loop.close()
-
-    except SystemExit as error:
-        # TODO: Remove this workaround once we can improve Jedi
-        # disconnection logic
-        tb = "".join(traceback.format_tb(error.__traceback__))
-        if tb.find("connection_lost") > 0:
-            logging.warning("Positron Language Server client disconnected, exiting.")
-            exitStatus = 0
-        else:
-            logging.error("Error in Positron Language Server: %s", error)
-            exitStatus = 1
-
-    sys.exit(exitStatus)
+    try:
+        loop.run_forever()
+    except (KeyboardInterrupt, SystemExit):
+        logger.exception("Unexpected exception in event loop")
+        exit_status = 1
+    finally:
+        loop.close()
+
+    logger.info(f"Exiting process with status {exit_status}")
+    sys.exit(exit_status)