refactor(langserver): correct refresh handling and remove some unneeded code

Author: d-biehl

packages/core/src/robotcode/core/async_itertools.py

@@ -519,29 +519,6 @@ def set_result(w: asyncio.Future[Any], ev: threading.Event) -> None:
             self._wake_up_first()
 
 
-class RLock(Lock):
-    def __init__(self) -> None:
-        super().__init__()
-        self._task: Optional[asyncio.Task[Any]] = None
-        self._depth = 0
-
-    async def acquire(self) -> bool:
-        if self._task is None or self._task != asyncio.current_task():
-            await super().acquire()
-            self._task = asyncio.current_task()
-            assert self._depth == 0
-        self._depth += 1
-
-        return True
-
-    def release(self) -> None:
-        if self._depth > 0:
-            self._depth -= 1
-        if self._depth == 0:
-            super().release()
-            self._task = None
-
-
 _global_futures_set: Set[asyncio.Future[Any]] = set()
 
 

packages/core/src/robotcode/core/async_tools.py

@@ -3,6 +3,8 @@
 from threading import Event, RLock, Thread, current_thread, local
 from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, TypeVar, cast, overload
 
+from typing_extensions import ParamSpec
+
 _F = TypeVar("_F", bound=Callable[..., Any])
 _TResult = TypeVar("_TResult")
 
@@ -118,7 +120,10 @@ def _remove_future_from_running_callables(future: FutureEx[Any]) -> None:
         _running_callables.pop(future, None)
 
 
-def run_in_thread(callable: Callable[..., _TResult], *args: Any, **kwargs: Any) -> FutureEx[_TResult]:
+_P = ParamSpec("_P")
+
+
+def run_in_thread(callable: Callable[_P, _TResult], *args: _P.args, **kwargs: _P.kwargs) -> FutureEx[_TResult]:
     future: FutureEx[_TResult] = FutureEx()
     with _running_callables_lock:
         thread = Thread(

packages/core/src/robotcode/core/concurrent.py

@@ -64,7 +64,7 @@ def __del__(self) -> None:
         self.close()
 
     @_logger.call
-    async def on_connection_lost(self, sender: Any, exc: Optional[BaseException]) -> None:
+    def on_connection_lost(self, sender: Any, exc: Optional[BaseException]) -> None:
         if sender == self._protocol:
             self._protocol = None
 

packages/debugger/src/robotcode/debugger/launcher/client.py

@@ -33,11 +33,11 @@
 )
 
 from robotcode.core.async_tools import (
-    async_event,
     create_sub_task,
     run_coroutine_in_thread,
 )
 from robotcode.core.concurrent import FutureEx, is_threaded_callable, run_in_thread
+from robotcode.core.event import event
 from robotcode.core.utils.dataclasses import as_json, from_dict
 from robotcode.core.utils.inspect import ensure_coroutine, iter_methods
 from robotcode.core.utils.logging import LoggingDescriptor
@@ -374,12 +374,12 @@ def __init__(self) -> None:
     def loop(self) -> Optional[asyncio.AbstractEventLoop]:
         return self._loop
 
-    @async_event
-    async def on_connection_made(sender, transport: asyncio.BaseTransport) -> None:
+    @event
+    def on_connection_made(sender, transport: asyncio.BaseTransport) -> None:
         ...
 
-    @async_event
-    async def on_connection_lost(sender, exc: Optional[BaseException]) -> None:
+    @event
+    def on_connection_lost(sender, exc: Optional[BaseException]) -> None:
         ...
 
     def connection_made(self, transport: asyncio.BaseTransport) -> None:
@@ -390,10 +390,10 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None:
         if isinstance(transport, asyncio.WriteTransport):
             self.write_transport = transport
 
-        create_sub_task(self.on_connection_made(self, transport))
+        self.on_connection_made(self, transport)
 
     def connection_lost(self, exc: Optional[BaseException]) -> None:
-        create_sub_task(self.on_connection_lost(self, exc))
+        self.on_connection_lost(self, exc)
         self._loop = None
 
     def eof_received(self) -> Optional[bool]:
@@ -451,6 +451,7 @@ def __init__(self) -> None:
         self._received_request: OrderedDict[Union[str, int, None], ReceivedRequestEntry] = OrderedDict()
         self._received_request_lock = threading.RLock()
         self._signature_cache: Dict[Callable[..., Any], inspect.Signature] = {}
+        self._running_handle_message_tasks: Set[asyncio.Future[Any]] = set()
 
     @staticmethod
     def _generate_json_rpc_messages_from_dict(
@@ -494,15 +495,10 @@ def _handle_body(self, body: bytes, charset: str) -> None:
             self.send_error(JsonRPCErrors.PARSE_ERROR, f"{type(e).__name__}: {e}")
 
     def _handle_messages(self, iterator: Iterator[JsonRPCMessage]) -> None:
-        def done(f: asyncio.Future[Any]) -> None:
-            if f.done() and not f.cancelled():
-                ex = f.exception()
-
-                if ex is None or isinstance(ex, asyncio.CancelledError):
-                    return
-
         for m in iterator:
-            create_sub_task(self.handle_message(m)).add_done_callback(done)
+            task = asyncio.create_task(self.handle_message(m))
+            self._running_handle_message_tasks.add(task)
+            task.add_done_callback(self._running_handle_message_tasks.discard)
 
     @__logger.call
     async def handle_message(self, message: JsonRPCMessage) -> None:

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

@@ -1,7 +1,7 @@
 from concurrent.futures import CancelledError
 from typing import TYPE_CHECKING, Any, Final, List, Optional
 
-from robotcode.core.concurrent import FutureEx, check_current_thread_canceled, threaded
+from robotcode.core.concurrent import FutureEx, check_current_thread_canceled, run_in_thread, threaded
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     CodeLens,
@@ -90,13 +90,20 @@ def _code_lens_resolve(self, params: CodeLens, *args: Any, **kwargs: Any) -> Cod
 
         return params
 
-    def refresh(self) -> None:
-        if not (
+    def refresh(self, now: bool = True) -> None:
+        if self.refresh_task is not None and not self.refresh_task.done():
+            self.refresh_task.cancel()
+
+        self.refresh_task = run_in_thread(self._refresh, now)
+
+    def _refresh(self, now: bool = True) -> None:
+        if (
             self.parent.client_capabilities is not None
             and self.parent.client_capabilities.workspace is not None
             and self.parent.client_capabilities.workspace.code_lens is not None
             and self.parent.client_capabilities.workspace.code_lens.refresh_support
         ):
-            return
+            if not now:
+                check_current_thread_canceled(1)
 
-        self.parent.send_request("workspace/codeLens/refresh").result(self._refresh_timeout)
+            self.parent.send_request("workspace/codeLens/refresh").result(self._refresh_timeout)

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

@@ -100,6 +100,8 @@ def __init__(self, protocol: "LanguageServerProtocol") -> None:
 
         self.client_supports_pull = False
 
+        self._refresh_timeout = 5
+
     def server_initialized(self, sender: Any) -> None:
         self._workspace_diagnostics_task = run_in_thread(self.run_workspace_diagnostics)
 
@@ -221,6 +223,8 @@ def run_workspace_diagnostics(self) -> None:
                     "Analyse workspace", cancellable=False, current=0, max=len(documents) + 1, start=False
                 ) as progress:
                     for i, document in enumerate(documents):
+                        check_current_thread_canceled()
+
                         mode = self.get_diagnostics_mode(document.uri)
                         if mode == DiagnosticsMode.OFF:
                             self.get_diagnostics_data(document).version = document.version
@@ -418,23 +422,20 @@ def get_diagnostics_mode(self, uri: Uri) -> DiagnosticsMode:
 
         return DiagnosticsMode.OPENFILESONLY
 
-    def __do_refresh(self, now: bool = False) -> None:
-        if not now:
-            check_current_thread_canceled(1)
-
-        self.__refresh()
-
     def refresh(self, now: bool = False) -> None:
         if self.refresh_task is not None and not self.refresh_task.done():
             self.refresh_task.cancel()
 
-        self.refresh_task = run_in_thread(self.__do_refresh, now)
+        self.refresh_task = run_in_thread(self._refresh, now)
 
-    def __refresh(self) -> None:
+    def _refresh(self, now: bool = False) -> None:
         if (
             self.parent.client_capabilities
             and self.parent.client_capabilities.workspace
             and self.parent.client_capabilities.workspace.diagnostics
             and self.parent.client_capabilities.workspace.diagnostics.refresh_support
         ):
-            self.parent.send_request("workspace/diagnostic/refresh").result(30)
+            if not now:
+                check_current_thread_canceled(1)
+
+            self.parent.send_request("workspace/diagnostic/refresh").result(self._refresh_timeout)

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

@@ -1,7 +1,7 @@
 from concurrent.futures import CancelledError
 from typing import TYPE_CHECKING, Any, Final, List, Optional
 
-from robotcode.core.concurrent import threaded
+from robotcode.core.concurrent import FutureEx, check_current_thread_canceled, run_in_thread, threaded
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     InlayHint,
@@ -27,6 +27,8 @@ class InlayHintProtocolPart(LanguageServerProtocolPart):
 
     def __init__(self, parent: "LanguageServerProtocol") -> None:
         super().__init__(parent)
+        self.refresh_task: Optional[FutureEx[Any]] = None
+        self._refresh_timeout = 5
 
     @event
     def collect(sender, document: TextDocument, range: Range) -> Optional[List[InlayHint]]:  # NOSONAR
@@ -98,11 +100,20 @@ def _inlay_hint_resolve(
 
         return params
 
-    def refresh(self) -> None:
+    def refresh(self, now: bool = True) -> None:
+        if self.refresh_task is not None and not self.refresh_task.done():
+            self.refresh_task.cancel()
+
+        self.refresh_task = run_in_thread(self._refresh, now)
+
+    def _refresh(self, now: bool = True) -> None:
         if (
             self.parent.client_capabilities is not None
             and self.parent.client_capabilities.workspace is not None
             and self.parent.client_capabilities.workspace.inlay_hint is not None
             and self.parent.client_capabilities.workspace.inlay_hint.refresh_support
         ):
-            self.parent.send_request("workspace/inlayHint/refresh").result(30)
+            if not now:
+                check_current_thread_canceled(1)
+
+            self.parent.send_request("workspace/inlayHint/refresh").result(self._refresh_timeout)

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

@@ -1,7 +1,7 @@
 from asyncio import CancelledError
 from typing import TYPE_CHECKING, Any, Final, List, Optional
 
-from robotcode.core.concurrent import check_current_thread_canceled, threaded
+from robotcode.core.concurrent import FutureEx, check_current_thread_canceled, run_in_thread, threaded
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     DocumentSelector,
@@ -31,6 +31,8 @@ class InlineValueProtocolPart(LanguageServerProtocolPart):
 
     def __init__(self, parent: "LanguageServerProtocol") -> None:
         super().__init__(parent)
+        self.refresh_task: Optional[FutureEx[Any]] = None
+        self._refresh_timeout = 5
 
     @event
     def collect(
@@ -88,11 +90,20 @@ def _text_document_inline_value(
 
         return results
 
-    def refresh(self) -> None:
+    def refresh(self, now: bool = True) -> None:
+        if self.refresh_task is not None and not self.refresh_task.done():
+            self.refresh_task.cancel()
+
+        self.refresh_task = run_in_thread(self._refresh, now)
+
+    def _refresh(self, now: bool = True) -> None:
         if (
             self.parent.client_capabilities
             and self.parent.client_capabilities.workspace
             and self.parent.client_capabilities.workspace.inline_value
             and self.parent.client_capabilities.workspace.inline_value.refresh_support
         ):
-            self.parent.send_request("workspace/inlineValue/refresh").result(30)
+            if not now:
+                check_current_thread_canceled(1)
+
+            self.parent.send_request("workspace/inlineValue/refresh").result(self._refresh_timeout)