Pygls handles requests that are known to be cancelled #517
Description
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)