Remove old synchronous code from Reader and Writer

Author: FlyingSamson

python_packages/jupyter_lsp/jupyter_lsp/connection.py

@@ -7,19 +7,17 @@
 > > MIT License   https://github.com/palantir/python-jsonrpc-server/blob/0.2.0/LICENSE
 > > Copyright 2018 Palantir Technologies, Inc.
 """
-import os
 from abc import ABC, ABCMeta, abstractmethod
-from typing import List, Optional, Text
+from typing import Optional, Text
 
 # pylint: disable=broad-except
 import anyio
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextSendStream
-from tornado.gen import convert_yielded
 from tornado.httputil import HTTPHeaders
 from tornado.ioloop import IOLoop
 from tornado.queues import Queue
-from traitlets import Float, Instance, Int, default
+from traitlets import Instance, Int
 from traitlets.config import LoggingConfigurable
 from traitlets.traitlets import MetaHasTraits
 
@@ -48,15 +46,8 @@ async def close(self):
 
 
 class LspStreamReader(LspStreamBase):
-    """Language Server Reader
+    """Language Server Reader"""
 
-    Because non-blocking (but still synchronous) IO is used, rudimentary
-    exponential backoff is used.
-    """
-
-    max_wait = Float(help="maximum time to wait on idle stream").tag(config=True)
-    min_wait = Float(0.05, help="minimum time to wait on idle stream").tag(config=True)
-    next_wait = Float(0.05, help="next time to wait on idle stream").tag(config=True)
     receive_max_bytes = Int(
         65536,
         help="the maximum size a header line send by the language server may have",
@@ -74,32 +65,12 @@ async def close(self):
         await self.stream.aclose()
         self.log.debug("%s closed", self)
 
-    @default("max_wait")
-    def _default_max_wait(self):
-        return 0.1 if os.name == "nt" else self.min_wait * 2
-
-    async def sleep(self):
-        """Simple exponential backoff for sleeping"""
-        self.next_wait = min(self.next_wait * 2, self.max_wait)
-        await anyio.sleep(self.next_wait)
-
-    def wake(self):
-        """Reset the wait time"""
-        self.wait = self.min_wait
-
     async def read(self) -> None:
         """Read from a Language Server until it is closed"""
         while True:
             message = None
             try:
                 message = await self.read_one()
-
-                if not message:
-                    await self.sleep()
-                    continue
-                else:
-                    self.wake()
-
                 IOLoop.current().add_callback(self.queue.put_nowait, message)
             except anyio.ClosedResourceError:
                 # stream was closed -> terminate
@@ -109,89 +80,48 @@ async def read(self) -> None:
                 self.log.exception(
                     "%s couldn't enqueue message: %s (%s)", self, message, e
                 )
-                await self.sleep()
-
-    async def _read_content(
-        self, length: int, max_parts=1000, max_empties=200
-    ) -> Optional[bytes]:
-        """Read the full length of the message unless exceeding max_parts or
-           max_empties empty reads occur.
 
-        See https://github.com/jupyter-lsp/jupyterlab-lsp/issues/450
-
-        Crucial docs or read():
-            "If the argument is positive, and the underlying raw
-             stream is not interactive, multiple raw reads may be issued
-             to satisfy the byte count (unless EOF is reached first)"
+    async def _read_content(self, length: int) -> Optional[bytes]:
+        """Read the full length of the message.
 
         Args:
            - length: the content length
-           - max_parts: prevent absurdly long messages (1000 parts is several MBs):
-             1 part is usually sufficient but not enough for some long
-             messages 2 or 3 parts are often needed.
         """
-        raw = None
-        raw_parts: List[bytes] = []
-        received_size = 0
-        while received_size < length and len(raw_parts) < max_parts and max_empties > 0:
-            part = None
-            try:
-                part = await self.stream.receive_exactly(length - received_size)
-            except anyio.IncompleteRead:  # pragma: no cover
-                pass
-            if part is None:  # pragma: no cover
-                max_empties -= 1
-                await self.sleep()
-                continue
-            received_size += len(part)
-            raw_parts.append(part)
-
-        if raw_parts:
-            raw = b"".join(raw_parts)
-            if len(raw) != length:  # pragma: no cover
-                self.log.warning(
-                    f"Readout and content-length mismatch: {len(raw)} vs {length};"
-                    f"remaining empties: {max_empties}; remaining parts: {max_parts}"
-                )
-
-        return raw
+        try:
+            return await self.stream.receive_exactly(length)
+        except anyio.IncompleteRead:  # pragma: no cover
+            # resource has been closed before the requested bytes could be retrieved
+            # -> signal recource closed
+            raise anyio.ClosedResourceError
 
     async def read_one(self) -> Text:
         """Read a single message"""
         message = ""
         headers = HTTPHeaders()
 
-        line = await convert_yielded(self._readline())
+        line = await self._readline()
 
         if line:
             while line and line.strip():
                 headers.parse_line(line)
-                line = await convert_yielded(self._readline())
+                line = await self._readline()
 
             content_length = int(headers.get("content-length", "0"))
 
             if content_length:
                 raw = await self._read_content(length=content_length)
-                if raw is not None:
-                    message = raw.decode("utf-8").strip()
-                else:  # pragma: no cover
-                    self.log.warning(
-                        "%s failed to read message of length %s",
-                        self,
-                        content_length,
-                    )
+                message = raw.decode("utf-8").strip()
 
         return message
 
     async def _readline(self) -> Text:
-        """Read a line (or immediately return None)"""
+        """Read a line"""
         try:
             # use same max_bytes as is default for receive for now. It seems there is no
             # way of getting the bytes read until max_bytes is reached, so we cannot
             # iterate the receive_until call with smaller max_bytes values
-            async with anyio.move_on_after(0.2):
-                line = await self.stream.receive_until(b"\r\n", self.receive_max_bytes)
-                return line.decode("utf-8").strip()
+            line = await self.stream.receive_until(b"\r\n", self.receive_max_bytes)
+            return line.decode("utf-8").strip()
         except anyio.IncompleteRead:
             # resource has been closed before the requested bytes could be retrieved
             # -> signal recource closed
@@ -225,7 +155,7 @@ async def write(self) -> None:
             try:
                 n_bytes = len(message.encode("utf-8"))
                 response = "Content-Length: {}\r\n\r\n{}".format(n_bytes, message)
-                await convert_yielded(self._write_one(response))
+                await self._write_one(response)
             except (
                 anyio.ClosedResourceError,
                 anyio.BrokenResourceError,

python_packages/jupyter_lsp/jupyter_lsp/session.py

@@ -57,6 +57,8 @@ class LanguageServerSessionBase(
     )
     main_loop = Instance(
         IOLoop, help="the event loop of the main thread", allow_none=True)
+    thread_loop = Instance(
+        IOLoop, help="the event loop of the worker thread", allow_none=True)
     writer = Instance(LspStreamWriter, help="the JSON-RPC writer", allow_none=True)
     reader = Instance(LspStreamReader, help="the JSON-RPC reader", allow_none=True)
     from_lsp = Instance(
@@ -118,8 +120,7 @@ def start(self):
     def stop(self):
         """shut down the session"""
         if self.cancelscope is not None:
-            self.cancelscope.cancel()
-            self.cancelscope = None
+            self.thread_loop.add_callback(self.cancelscope.cancel)
 
         # wait for the session to get cleaned up
         if self.thread and self.thread.is_alive():
@@ -132,12 +133,15 @@ async def run(self):
         the event `self.started` will be set when everything is set up and the session
         will be ready for communication
         """
+        self.thread_loop = IOLoop.current()
         async with CancelScope() as scope:
             self.cancelscope = scope
             await self.initialize()
             self.started.set()
             await self.listen()
         await self.cleanup()
+        self.cancelscope = None
+        self.thread_loop = None
 
     async def initialize(self):
         """initialize a language server session"""
@@ -187,7 +191,7 @@ def _on_handlers(self, change: Bunch):
     def write(self, message):
         """wrapper around the write queue to keep it mostly internal"""
         self.last_handler_message_at = self.now()
-        IOLoop.current().add_callback(self.to_lsp.put_nowait, message)
+        self.thread_loop.add_callback(self.to_lsp.put_nowait, message)
 
     def now(self):
         return datetime.now(timezone.utc)