Pygls handles requests that are known to be cancelled

[seeM]

Firstly, hello and thank you for your awesome work on pygls! 🥳 We've been using it (via jedi-language-server) over at Positron.

One of our users recently encountered an interesting pygls performance issue where a very slow request causes a buildup of cancelled requests that still get processed unnecessarily by the server. It would be great if pygls could skip requests that are cancelled before it starts handling them.

Details

The client, Positron in this case, starts with a textDocument/hover request.

The client times out waiting for the server to respond to the request, sends a $/cancelRequest, and tries a new textDocument/hover. It may do this multiple times all while the first request is still being handled. If the handler is synchronous (as are most jedi-language-server handlers) and the request takes very long (30 seconds in our user's case) it's possible for a pretty large queue to build up (~70 messages in our case), many of which are cancellations of the other messages (~30 cancellations in our case, meaning that ~60/70 messages needn't be handled).

However, pygls handles each message and then handles its cancellation (leading to warning logs Cancel notification for unknown message id). It could take very long for all of these cancelled messages to be handled, all while the user sees no response in the UI because the client thinks they've cancelled everything (the buildup wasn't resolved after ~90 additional seconds in our case).

I think this is worth optimizing in pygls. If there's also a way for us to configure our handlers to avoid this situation entirely, any pointers would be appreciated!

Workaround

I intend on overriding LanguageServerProtocol as follows as a temporary workaround:

    def _data_received(self, data: bytes) -> None:
        self._messages_to_handle = []

        # Read the received bytes and call `self._procedure_handler` with each message,
        # which should actually handle each message but we've overridden to just add them to a queue.
        super()._data_received(data)

        def is_request(message):
            return hasattr(message, "method") and hasattr(message, "id")

        def is_cancel_notification(message):
            return getattr(message, "method", None) == CANCEL_REQUEST

        # First pass: find all requests that were cancelled in the same batch of `data`.
        request_ids = set()
        cancelled_ids = set()
        for message in self._messages_to_handle:
            if is_request(message):
                request_ids.add(message.id)
            elif is_cancel_notification(message) and message.params.id in request_ids:
                cancelled_ids.add(message.params.id)

        # Second pass: remove all requests that were cancelled in the same batch of `data`,
        # and the cancel notifications themselves.
        self._messages_to_handle = [
            msg
            for msg in self._messages_to_handle
            if not (
                # Remove cancel notifications whose params.id is in cancelled_ids...
                (is_cancel_notification(msg) and msg.params.id in cancelled_ids)
                # ...or original messages whose id is in cancelled_ids.
                or (is_request(msg) and msg.id in cancelled_ids)
            )
        ]

        # Now handle the messages.
        for message in self._messages_to_handle:
            super()._procedure_handler(message)

    def _procedure_handler(self, message) -> None:
        self._messages_to_handle.append(message)

[tombh]

Thank you for using Pygls! 🥳

So you're actively using that workaround at the moment with success?

We're actually in the middle of open alpha for migrating to v2 of Pygls. Have you tried it? It's on the main branch, we have a migration guide. It has some notable changes in the area you're talking about.

We'd be very much open to PRs, either targeting the current release or current v2 alpha. I'm not familiar with the new code so it might be some time before I could be in a position to contribute a fix.

[seeM]

@tombh I have a PR in our repo (posit-dev/positron#6219) and a unit test but haven't validated it end to end.

I haven't tried v2 yet. We had to make a few asyncio/threading changes to LanguageServer on v1 to get things working well in our context. I assume it'll need some fiddling to get that all working on v2 and I haven't had a chance to dive into that yet. We also intend on keeping support for Python 3.8 for a bit longer.

I unfortunately can't promise a PR right now, but might be able to look into it when we upgrade to v2.

[alcarney]

If there's also a way for us to configure our handlers to avoid this situation entirely, any pointers would be appreciated!

The main thing to try would be trying to offload the heavy computation from the main thread so that the main loop has chance to reach the cancel notifications and process them...

I see from your linked PR that you found the @server.thread() decorator - does that make a difference without applying your patch?

You could also try an async handler wrapping a subprocess, (which appears to work well for me in esbonio), but that of course takes a lot more to setup.

Something to note, when running the server with start_io I'm fairly sure that the data given to _data_received will only ever have one message in it - not sure if that affects your workaround at all? _data_received is called from aio_readline in case that helps at all

[seeM]

@alcarney Thank you for the tips!

Using a subprocess won't be easy since we'd need to read some data across processes, but we'll try @server.thread() and see how that goes. The downside of thread() is that it isn't cancellable (I think?), but it should circumvent the underlying issue to begin with.