refactor(langserver): make documents handling synchronous and more threadsafe

Author: d-biehl

packages/core/src/robotcode/core/event.py

@@ -69,19 +69,37 @@ def __iter__(self) -> Iterator[Callable[_TParams, _TResult]]:
             if c is not None:
                 yield c
 
-    def _notify(self, *__args: _TParams.args, **__kwargs: _TParams.kwargs) -> Iterator[_TResult]:
-        for method in set(self):
+    def _notify(
+        self,
+        *__args: _TParams.args,
+        callback_filter: Optional[Callable[[Callable[..., Any]], bool]] = None,
+        **__kwargs: _TParams.kwargs,
+    ) -> Iterator[_TResult]:
+        for method in filter(
+            lambda x: callback_filter(x) if callback_filter is not None else True,
+            set(self),
+        ):
             yield method(*__args, **__kwargs)
 
 
 class EventIterator(EventResultIteratorBase[_TParams, _TResult]):
-    def __call__(self, *__args: _TParams.args, **__kwargs: _TParams.kwargs) -> Iterator[_TResult]:
-        return self._notify(*__args, **__kwargs)
+    def __call__(
+        self,
+        *__args: _TParams.args,
+        callback_filter: Optional[Callable[[Callable[..., Any]], bool]] = None,
+        **__kwargs: _TParams.kwargs,
+    ) -> Iterator[_TResult]:
+        return self._notify(*__args, callback_filter=callback_filter, **__kwargs)
 
 
 class Event(EventResultIteratorBase[_TParams, _TResult]):
-    def __call__(self, *__args: _TParams.args, **__kwargs: _TParams.kwargs) -> List[_TResult]:
-        return list(self._notify(*__args, **__kwargs))
+    def __call__(
+        self,
+        *__args: _TParams.args,
+        callback_filter: Optional[Callable[[Callable[..., Any]], bool]] = None,
+        **__kwargs: _TParams.kwargs,
+    ) -> List[_TResult]:
+        return list(self._notify(*__args, callback_filter=callback_filter, **__kwargs))
 
 
 _TEvent = TypeVar("_TEvent")
@@ -125,31 +143,3 @@ def __init__(self, _func: Callable[_TParams, _TResult]) -> None:
 class event(EventDescriptorBase[_TParams, _TResult, Event[_TParams, _TResult]]):  # noqa: N801
     def __init__(self, _func: Callable[_TParams, _TResult]) -> None:
         super().__init__(_func, Event[_TParams, _TResult])
-
-
-class ThreadedEventResultIteratorBase(EventResultIteratorBase[_TParams, _TResult]):
-    def _notify(self, *__args: _TParams.args, **__kwargs: _TParams.kwargs) -> Iterator[_TResult]:
-        for method in set(self):
-            yield method(*__args, **__kwargs)
-
-
-class ThreadedEventIterator(ThreadedEventResultIteratorBase[_TParams, _TResult]):
-    def __call__(self, *__args: _TParams.args, **__kwargs: _TParams.kwargs) -> Iterator[_TResult]:
-        return self._notify(*__args, **__kwargs)
-
-
-class ThreadedEvent(ThreadedEventResultIteratorBase[_TParams, _TResult]):
-    def __call__(self, *__args: _TParams.args, **__kwargs: _TParams.kwargs) -> List[_TResult]:
-        return list(self._notify(*__args, **__kwargs))
-
-
-class threaded_event_iterator(  # noqa: N801
-    EventDescriptorBase[_TParams, _TResult, ThreadedEventIterator[_TParams, _TResult]]
-):
-    def __init__(self, _func: Callable[_TParams, _TResult]) -> None:
-        super().__init__(_func, ThreadedEventIterator[_TParams, _TResult])
-
-
-class threaded_event(EventDescriptorBase[_TParams, _TResult, ThreadedEvent[_TParams, _TResult]]):  # noqa: N801
-    def __init__(self, _func: Callable[_TParams, _TResult]) -> None:
-        super().__init__(_func, ThreadedEvent[_TParams, _TResult])

packages/jsonrpc2/src/robotcode/jsonrpc2/protocol.py

@@ -144,12 +144,21 @@ class InvalidProtocolVersionError(JsonRPCParseError):
     pass
 
 
-class RpcMethodEntry(NamedTuple):
+@dataclass
+class RpcMethodEntry:
     name: str
     method: Callable[..., Any]
     param_type: Optional[Type[Any]]
     cancelable: bool
 
+    _is_coroutine: Optional[bool] = field(default=None, init=False)
+
+    @property
+    def is_coroutine(self) -> bool:
+        if self._is_coroutine is None:
+            self._is_coroutine = inspect.iscoroutinefunction(self.method)
+        return self._is_coroutine
+
 
 @runtime_checkable
 class RpcMethod(Protocol):
@@ -622,7 +631,7 @@ async def handle_error(self, message: JsonRPCError) -> None:
                     JsonRPCErrorException(message.error.code, message.error.message, message.error.data)
                 )
             else:
-                self.__logger.warning(lambda: f"Response for {message} is already done.")
+                self.__logger.warning(lambda: f"Error Response for {message} is already done.")
 
         except (SystemExit, KeyboardInterrupt):
             raise
@@ -707,53 +716,57 @@ async def handle_request(self, message: JsonRPCRequest) -> None:
 
         params = self._convert_params(e.method, e.param_type, message.params)
 
-        if (
-            isinstance(e.method, HasThreaded)
-            and e.method.__threaded__
-            or inspect.ismethod(e.method)
-            and isinstance(e.method.__func__, HasThreaded)
-            and e.method.__func__.__threaded__
-        ):
-            task = run_coroutine_in_thread(
-                ensure_coroutine(cast(Callable[..., Any], e.method)), *params[0], **params[1]
-            )
+        if not e.is_coroutine:
+            self.send_response(message.id, e.method(*params[0], **params[1]))
         else:
-            task = create_sub_task(
-                ensure_coroutine(e.method)(*params[0], **params[1]),
-                name=message.method,
-            )
+            if (
+                isinstance(e.method, HasThreaded)
+                and e.method.__threaded__
+                or inspect.ismethod(e.method)
+                and isinstance(e.method.__func__, HasThreaded)
+                and e.method.__func__.__threaded__
+            ):
+                task = run_coroutine_in_thread(
+                    ensure_coroutine(cast(Callable[..., Any], e.method)), *params[0], **params[1]
+                )
+            else:
+                task = create_sub_task(
+                    ensure_coroutine(e.method)(*params[0], **params[1]),
+                    name=message.method,
+                )
 
-        with self._received_request_lock:
-            self._received_request[message.id] = ReceivedRequestEntry(task, message, e.cancelable)
-
-        def done(t: asyncio.Future[Any]) -> None:
-            try:
-                if not t.cancelled():
-                    ex = t.exception()
-                    if ex is not None:
-                        self.__logger.exception(ex, exc_info=ex)
-                        raise JsonRPCErrorException(JsonRPCErrors.INTERNAL_ERROR, f"{type(ex).__name__}: {ex}") from ex
-
-                self.send_response(message.id, t.result())
-            except asyncio.CancelledError:
-                self.__logger.debug(lambda: f"request message {message!r} canceled")
-                self.send_error(JsonRPCErrors.REQUEST_CANCELLED, "Request canceled.", id=message.id)
-            except (SystemExit, KeyboardInterrupt):
-                raise
-            except JsonRPCErrorException as e:
-                self.send_error(e.code, e.message or f"{type(e).__name__}: {e}", id=message.id, data=e.data)
-            except BaseException as e:
-                self.__logger.exception(e)
-                self.send_error(JsonRPCErrors.INTERNAL_ERROR, f"{type(e).__name__}: {e}", id=message.id)
-            finally:
-                with self._received_request_lock:
-                    self._received_request.pop(message.id, None)
-
-        task.add_done_callback(done)
-
-        await task
+            with self._received_request_lock:
+                self._received_request[message.id] = ReceivedRequestEntry(task, message, e.cancelable)
+
+            def done(t: asyncio.Future[Any]) -> None:
+                try:
+                    if not t.cancelled():
+                        ex = t.exception()
+                        if ex is not None:
+                            self.__logger.exception(ex, exc_info=ex)
+                            raise JsonRPCErrorException(
+                                JsonRPCErrors.INTERNAL_ERROR, f"{type(ex).__name__}: {ex}"
+                            ) from ex
+
+                    self.send_response(message.id, t.result())
+                except asyncio.CancelledError:
+                    self.__logger.debug(lambda: f"request message {message!r} canceled")
+                    self.send_error(JsonRPCErrors.REQUEST_CANCELLED, "Request canceled.", id=message.id)
+                except (SystemExit, KeyboardInterrupt):
+                    raise
+                except JsonRPCErrorException as e:
+                    self.send_error(e.code, e.message or f"{type(e).__name__}: {e}", id=message.id, data=e.data)
+                except BaseException as e:
+                    self.__logger.exception(e)
+                    self.send_error(JsonRPCErrors.INTERNAL_ERROR, f"{type(e).__name__}: {e}", id=message.id)
+                finally:
+                    with self._received_request_lock:
+                        self._received_request.pop(message.id, None)
+
+            task.add_done_callback(done)
+
+            await task
 
-    @__logger.call
     def cancel_request(self, id: Union[int, str, None]) -> None:
         with self._received_request_lock:
             entry = self._received_request.get(id, None)
@@ -762,8 +775,7 @@ def cancel_request(self, id: Union[int, str, None]) -> None:
             self.__logger.debug(lambda: f"try to cancel request {entry.request if entry is not None else ''}")
             entry.future.cancel()
 
-    @__logger.call
-    async def cancel_all_received_request(self) -> None:
+    def cancel_all_received_request(self) -> None:
         for entry in self._received_request.values():
             if entry is not None and entry.cancelable and entry.future is not None and not entry.future.cancelled():
                 entry.future.cancel()
@@ -778,23 +790,26 @@ async def handle_notification(self, message: JsonRPCNotification) -> None:
         try:
             params = self._convert_params(e.method, e.param_type, message.params)
 
-            if (
-                isinstance(e.method, HasThreaded)
-                and e.method.__threaded__
-                or inspect.ismethod(e.method)
-                and isinstance(e.method.__func__, HasThreaded)
-                and e.method.__func__.__threaded__
-            ):
-                task = run_coroutine_in_thread(
-                    ensure_coroutine(cast(Callable[..., Any], e.method)), *params[0], **params[1]
-                )
+            if not e.is_coroutine:
+                e.method(*params[0], **params[1])
             else:
-                task = create_sub_task(
-                    ensure_coroutine(e.method)(*params[0], **params[1]),
-                    name=message.method,
-                )
+                if (
+                    isinstance(e.method, HasThreaded)
+                    and e.method.__threaded__
+                    or inspect.ismethod(e.method)
+                    and isinstance(e.method.__func__, HasThreaded)
+                    and e.method.__func__.__threaded__
+                ):
+                    task = run_coroutine_in_thread(
+                        ensure_coroutine(cast(Callable[..., Any], e.method)), *params[0], **params[1]
+                    )
+                else:
+                    task = create_sub_task(
+                        ensure_coroutine(e.method)(*params[0], **params[1]),
+                        name=message.method,
+                    )
 
-            await task
+                await task
         except asyncio.CancelledError:
             pass
         except (SystemExit, KeyboardInterrupt):

packages/language_server/src/robotcode/language_server/common/parts/code_action.py

@@ -71,7 +71,7 @@ async def _text_document_code_action(
     ) -> Optional[List[Union[Command, CodeAction]]]:
         results: List[Union[Command, CodeAction]] = []
 
-        document = await self.parent.documents.get(text_document.uri)
+        document = self.parent.documents.get(text_document.uri)
         if document is None:
             return None
 

packages/language_server/src/robotcode/language_server/common/parts/code_lens.py

@@ -49,7 +49,7 @@ async def _text_document_code_lens(
         self, text_document: TextDocumentIdentifier, *args: Any, **kwargs: Any
     ) -> Optional[List[CodeLens]]:
         results: List[CodeLens] = []
-        document = await self.parent.documents.get(text_document.uri)
+        document = self.parent.documents.get(text_document.uri)
         if document is None:
             return None
 

packages/language_server/src/robotcode/language_server/common/parts/completion.py

@@ -97,7 +97,7 @@ async def _text_document_completion(
         if context is not None and context.trigger_kind == CompletionTriggerKind.TRIGGER_CHARACTER:
             await asyncio.sleep(0.25)
 
-        document = await self.parent.documents.get(text_document.uri)
+        document = self.parent.documents.get(text_document.uri)
         if document is None:
             return None
 

packages/language_server/src/robotcode/language_server/common/parts/declaration.py

@@ -61,7 +61,7 @@ async def _text_document_declaration(
         locations: List[Location] = []
         location_links: List[LocationLink] = []
 
-        document = await self.parent.documents.get(text_document.uri)
+        document = self.parent.documents.get(text_document.uri)
         if document is None:
             return None
 
@@ -85,13 +85,13 @@ async def _text_document_declaration(
 
         if locations:
             for location in locations:
-                doc = await self.parent.documents.get(location.uri)
+                doc = self.parent.documents.get(location.uri)
                 if doc is not None:
                     location.range = doc.range_to_utf16(location.range)
 
         if location_links:
             for location_link in location_links:
-                doc = await self.parent.documents.get(location_link.target_uri)
+                doc = self.parent.documents.get(location_link.target_uri)
                 if doc is not None:
                     location_link.target_range = doc.range_to_utf16(location_link.target_range)
                     location_link.target_selection_range = doc.range_to_utf16(location_link.target_selection_range)

packages/language_server/src/robotcode/language_server/common/parts/definition.py

@@ -55,7 +55,7 @@ async def _text_document_definition(
         locations: List[Location] = []
         location_links: List[LocationLink] = []
 
-        document = await self.parent.documents.get(text_document.uri)
+        document = self.parent.documents.get(text_document.uri)
         if document is None:
             return None
 
@@ -81,13 +81,13 @@ async def _text_document_definition(
 
         if locations:
             for location in locations:
-                doc = await self.parent.documents.get(location.uri)
+                doc = self.parent.documents.get(location.uri)
                 if doc is not None:
                     location.range = doc.range_to_utf16(location.range)
 
         if location_links:
             for location_link in location_links:
-                doc = await self.parent.documents.get(location_link.target_uri)
+                doc = self.parent.documents.get(location_link.target_uri)
                 if doc is not None:
                     location_link.target_range = doc.range_to_utf16(location_link.target_range)
                     location_link.target_selection_range = doc.range_to_utf16(location_link.target_selection_range)

packages/language_server/src/robotcode/language_server/common/parts/diagnostics.py

@@ -281,7 +281,10 @@ async def force_refresh_document(self, document: TextDocument, refresh: bool = T
 
     @_logger.call
     @threaded()
-    async def on_did_close(self, sender: Any, document: TextDocument) -> None:
+    def on_did_close(self, sender: Any, document: TextDocument) -> None:
+        create_sub_task(self._close_diagnostics_for_document(document), loop=self.diagnostics_loop)
+
+    async def _close_diagnostics_for_document(self, document: TextDocument) -> None:
         if await self.get_diagnostics_mode(document.uri) == DiagnosticsMode.WORKSPACE:
             return
 
@@ -464,7 +467,7 @@ async def _get_diagnostics_for_document(
                         d.range = document.range_to_utf16(d.range)
 
                         for r in d.related_information or []:
-                            doc = await self.parent.documents.get(r.location.uri)
+                            doc = self.parent.documents.get(r.location.uri)
                             if doc is not None:
                                 r.location.range = doc.range_to_utf16(r.location.range)
 
@@ -494,7 +497,7 @@ def publish_diagnostics(self, document: TextDocument, diagnostics: List[Diagnost
             ),
         )
 
-    async def update_document_diagnostics(self, sender: Any, document: TextDocument) -> None:
+    def update_document_diagnostics(self, sender: Any, document: TextDocument) -> None:
         self.create_document_diagnostics_task(document, True)
 
     @rpc_method(name="textDocument/diagnostic", param_type=DocumentDiagnosticParams)
@@ -513,7 +516,7 @@ async def _text_document_diagnostic(
                     LSPErrorCodes.SERVER_CANCELLED, "Server not initialized.", DiagnosticServerCancellationData(True)
                 )
 
-            document = await self.parent.documents.get(text_document.uri)
+            document = self.parent.documents.get(text_document.uri)
             if document is None:
                 raise JsonRPCErrorException(LSPErrorCodes.SERVER_CANCELLED, f"Document {text_document!r} not found.")
 

packages/language_server/src/robotcode/language_server/common/parts/document_highlight.py

@@ -50,7 +50,7 @@ async def _text_document_document_highlight(
     ) -> Optional[List[DocumentHighlight]]:
         highlights: List[DocumentHighlight] = []
 
-        document = await self.parent.documents.get(text_document.uri)
+        document = self.parent.documents.get(text_document.uri)
         if document is None:
             return None