Switch from Tornado Queues to anyio MemoryObjectStreams

Author: FlyingSamson

python_packages/jupyter_lsp/jupyter_lsp/connection.py

@@ -14,9 +14,8 @@
 import anyio
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextSendStream
+from anyio.streams.stapled import StapledObjectStream
 from tornado.httputil import HTTPHeaders
-from tornado.ioloop import IOLoop
-from tornado.queues import Queue
 from traitlets import Instance, Int
 from traitlets.config import LoggingConfigurable
 from traitlets.traitlets import MetaHasTraits
@@ -31,7 +30,7 @@ class LspStreamBase(LoggingConfigurable, ABC, metaclass=LspStreamMeta):
     streams
     """
 
-    queue = Instance(Queue, help="queue to get/put")
+    queue = Instance(StapledObjectStream, help="queue to get/put")
 
     def __repr__(self):  # pragma: no cover
         return "<{}(parent={})>".format(self.__class__.__name__, self.parent)
@@ -71,7 +70,7 @@ async def read(self) -> None:
             message = None
             try:
                 message = await self.read_one()
-                IOLoop.current().add_callback(self.queue.put_nowait, message)
+                await self.queue.send(message)
             except anyio.ClosedResourceError:
                 # stream was closed -> terminate
                 self.log.debug("Stream closed while a read was still in progress")
@@ -151,7 +150,7 @@ async def close(self):
     async def write(self) -> None:
         """Write to a Language Server until it closes"""
         while True:
-            message = await self.queue.get()
+            message = await self.queue.receive()
             try:
                 n_bytes = len(message.encode("utf-8"))
                 response = "Content-Length: {}\r\n\r\n{}".format(n_bytes, message)
@@ -165,8 +164,6 @@ async def write(self) -> None:
                 break
             except Exception:  # pragma: no cover
                 self.log.exception("%s couldn't write message: %s", self, response)
-            finally:
-                self.queue.task_done()
 
     async def _write_one(self, message) -> None:
         await self.stream.send(message)

python_packages/jupyter_lsp/jupyter_lsp/session.py

@@ -1,6 +1,7 @@
 """ A session for managing a language server process
 """
 import atexit
+import math
 import os
 import string
 import subprocess
@@ -14,8 +15,8 @@
 import anyio
 from anyio import CancelScope
 from anyio.abc import Process, SocketStream
+from anyio.streams.stapled import StapledObjectStream
 from tornado.ioloop import IOLoop
-from tornado.queues import Queue
 from tornado.websocket import WebSocketHandler
 from traitlets import Bunch, Float, Instance, Set, Unicode, UseEnum, observe
 from traitlets.config import LoggingConfigurable
@@ -62,10 +63,14 @@ class LanguageServerSessionBase(
     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(
-        Queue, help="a queue for string messages from the server", allow_none=True
+        StapledObjectStream,
+        help="a queue for string messages from the server",
+        allow_none=True
     )
     to_lsp = Instance(
-        Queue, help="a queue for string message to the server", allow_none=True
+        StapledObjectStream,
+        help="a queue for string messages to the server",
+        allow_none=True
     )
     handlers = Set(
         trait=Instance(WebSocketHandler),
@@ -80,6 +85,10 @@ class LanguageServerSessionBase(
         5,
         help="timeout after which a process will be terminated forcefully",
     ).tag(config=True)
+    queue_size = Float(
+        math.inf,
+        help="the maximum number of messages that can be buffered in the queue"
+    ).tag(config=True)
 
     _skip_serialize = ["argv", "debug_argv"]
 
@@ -170,6 +179,12 @@ async def cleanup(self):
         if self.process is not None:
             await self.stop_process(self.stop_timeout)
             self.process = None
+        if self.from_lsp is not None:
+            await self.from_lsp.aclose()
+            self.from_lsp = None
+        if self.to_lsp is not None:
+            await self.to_lsp.aclose()
+            self.to_lsp = None
 
         self.status = SessionStatus.STOPPED
 
@@ -184,7 +199,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()
-        self.thread_loop.add_callback(self.to_lsp.put_nowait, message)
+        self.thread_loop.add_callback(self.to_lsp.send, message)
 
     def now(self):
         return datetime.now(timezone.utc)
@@ -242,8 +257,10 @@ async def stop_process(self, timeout: int = 5):
 
     def init_queues(self):
         """create the queues"""
-        self.from_lsp = Queue()
-        self.to_lsp = Queue()
+        self.from_lsp = StapledObjectStream(
+            *anyio.create_memory_object_stream(max_buffer_size=self.queue_size))
+        self.to_lsp = StapledObjectStream(
+            *anyio.create_memory_object_stream(max_buffer_size=self.queue_size))
 
     def substitute_env(self, env, base):
         final_env = copy(os.environ)
@@ -286,7 +303,6 @@ async def _broadcast_from_lsp(self):
             self.last_server_message_at = self.now()
             # handle message in the main thread's event loop
             self.main_loop.add_callback(self.parent.on_server_message, message, self)
-            self.from_lsp.task_done()
 
 
 class LanguageServerSessionStdio(LanguageServerSessionBase):

python_packages/jupyter_lsp/jupyter_lsp/tests/test_reader.py

@@ -1,8 +1,9 @@
 import subprocess
 
 import anyio
+from anyio.streams.stapled import StapledObjectStream
+import math
 import pytest
-from tornado.queues import Queue
 
 from jupyter_lsp.connection import LspStreamReader
 from jupyter_lsp.utils import get_unused_port
@@ -127,7 +128,8 @@ async def join_process(process: anyio.abc.Process, headstart=1, timeout=1):
 async def test_reader(
     message, repeats, interval, add_excess, mode, communicator_spawner
 ):
-    queue = Queue()
+    queue = StapledObjectStream(
+        *anyio.create_memory_object_stream(max_buffer_size=math.inf))
 
     port = get_unused_port() if mode == "tcp" else None
     process = await communicator_spawner.spawn_writer(
@@ -154,5 +156,5 @@ async def test_reader(
     if port is not None:
         await stream.aclose()
 
-    result = queue.get_nowait()
+    result = await queue.receive()
     assert result == message * repeats