Send websocket payload using a queue

[Raekkeri]

Hi,
we experienced JSON responses from the websocket server being mixed up if the requests from client (code editor, Monaco in our case) are too frequent. The response payload was basically unparseable JSON, and it seemed like a mashed up content of multiple JSON responses. Handling the send operations one-by-one by using the asyncio.Queue seemed to fix the issue, so here is a suggestion to apply the same fix into the main codebase.

[krassowski]

It looks like it needs a pass with ruff

[nneskildsf]

My project which uses python-lsp-server has been seriously haunted by this bug... We couldn't figure out what on earth was going on, but this makes total sense. Thank you for contributing this fix @Raekkeri.

[krassowski]

The response payload was basically unparseable JSON, and it seemed like a mashed up content of multiple JSON responses

Just leaving a note that I saw this happen before too. I recall it required many requests to be made in a very short period, and only occurred with specific setup of IO (which could be that it was just making it fast enough to allow this issue to occur).

I wonder if there is a way to reliably reproduce this issue and add a test which would confirm that this patch solves it.

[krassowski]

I am working on a test for this. For now this patch does not work on Python 3.9, it prevents the server from starting up at all in websocket mode:

$ pylsp --ws --port 6000
Traceback (most recent call last):
  File "python-lsp-server/bin/pylsp", line 33, in <module>
    sys.exit(load_entry_point('python-lsp-server', 'console_scripts', 'pylsp')())
  File "python-lsp-server/pylsp/__main__.py", line 81, in main
    start_ws_lang_server(args.port, args.check_parent_process, PythonLSPServer)
  File "python-lsp-server/pylsp/python_lsp.py", line 161, in start_ws_lang_server
    asyncio.run(run_server())
  File "python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "python3.9/asyncio/base_events.py", line 642, in run_until_complete
    return future.result()
  File "python-lsp-server/pylsp/python_lsp.py", line 158, in run_server
    payload, websocket = await send_queue.get()
  File "python3.9/asyncio/queues.py", line 166, in get
    await getter
RuntimeError: Task <Task pending name='Task-1' coro=<start_ws_lang_server.<locals>.run_server() running at python-lsp-server/pylsp/python_lsp.py:158> cb=[_run_until_complete_cb() at python3.9/asyncio/base_events.py:184]> got Future <Future pending> attached to a different loop

[krassowski]

So I have a test on krassowski#1 which fails when -vv is set (and pipe is used) but passes when -vv is absent (on main). I am yet to make it fail on main. This might be completely unrelated, but just noting down this oddity (bug) with -vv.

[krassowski]

The following diff makes this PR run fine for me:

diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index 2826ea1..e2f98c4 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -117,7 +117,8 @@ def start_ws_lang_server(port, check_parent_process, handler_class) -> None:
         ) from e

     with ThreadPoolExecutor(max_workers=10) as tpool:
-        send_queue = asyncio.Queue()
+        send_queue = None
+        loop = None

         async def pylsp_ws(websocket):
             log.debug("Creating LSP object")
@@ -147,11 +148,15 @@ def start_ws_lang_server(port, check_parent_process, handler_class) -> None:
             """Handler to send responses of  processed requests to respective web socket clients"""
             try:
                 payload = json.dumps(message, ensure_ascii=False)
-                asyncio.run(send_queue.put((payload, websocket)))
+                loop.call_soon_threadsafe(send_queue.put_nowait, (payload, websocket))
             except Exception as e:
                 log.exception("Failed to write message %s, %s", message, str(e))

         async def run_server():
+            nonlocal send_queue, loop
+            send_queue = asyncio.Queue()
+            loop = asyncio.get_running_loop()
+
             async with websockets.serve(pylsp_ws, port=port):
                 while 1:
                     # Wait until payload is available for sending

Still no luck in creating an isolated reproduction. Any ideas?

[krassowski]

I added a simple end-to-end test for websocket server running and returning replies to make sure we do not introduce the RuntimeError issue. It fails as seen here which confirms that we should not merge the PR in the original shape.

I was not able to create a test which would reproduce the malformed JSON issue, so this is not a test for the motivating problem, but the test that I added includes a best effort to simulate conditions which would usually lead to it, so potentially someone else could pick it up in another PR.

I will now push the patch from #633 (comment) to fix the failing test.

[krassowski]

@Raekkeri @nneskildsf can you confirm that the patch after my changes still solves the issue for you? Was the original one running well for you (it was failing for me locally and on CI once end-to-end test was added)?