refactor(langserver): remove threaded decorator from rpc methods and …

…introduce a threaded property for the rpc_method decorator
robotcodedev · Jan 5, 2024 · b478ae3 · b478ae3
1 parent 660d8f9
commit b478ae3
Show file tree

Hide file tree

Showing 27 changed files with 86 additions and 172 deletions.
diff --git a/packages/jsonrpc2/src/robotcode/jsonrpc2/protocol.py b/packages/jsonrpc2/src/robotcode/jsonrpc2/protocol.py
@@ -32,12 +32,8 @@
     runtime_checkable,
 )
 
-from robotcode.core.async_tools import create_sub_task, run_coroutine_in_thread
-from robotcode.core.concurrent import (
-    FutureEx,
-    is_threaded_callable,
-    run_in_thread,
-)
+from robotcode.core.async_tools import run_coroutine_in_thread
+from robotcode.core.concurrent import FutureEx, 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
@@ -150,6 +146,7 @@ class RpcMethodEntry:
     method: Callable[..., Any]
     param_type: Optional[Type[Any]]
     cancelable: bool
+    threaded: bool
 
     _is_coroutine: Optional[bool] = field(default=None, init=False)
 
@@ -181,6 +178,7 @@ def rpc_method(
     name: Optional[str] = None,
     param_type: Optional[Type[Any]] = None,
     cancelable: bool = True,
+    threaded: bool = False,
 ) -> Callable[[_F], _F]:
     ...
 
@@ -191,6 +189,7 @@ def rpc_method(
     name: Optional[str] = None,
     param_type: Optional[Type[Any]] = None,
     cancelable: bool = True,
+    threaded: bool = False,
 ) -> Callable[[_F], _F]:
     def _decorator(func: _F) -> Callable[[_F], _F]:
         if inspect.isclass(_func):
@@ -207,7 +206,7 @@ def _decorator(func: _F) -> Callable[[_F], _F]:
         if real_name is None or not real_name:
             raise ValueError("name is empty.")
 
-        cast(RpcMethod, f).__rpc_method__ = RpcMethodEntry(real_name, f, param_type, cancelable)
+        cast(RpcMethod, f).__rpc_method__ = RpcMethodEntry(real_name, f, param_type, cancelable, threaded)
         return func
 
     if _func is None:
@@ -274,6 +273,7 @@ def get_methods(obj: Any) -> Dict[str, RpcMethodEntry]:
                     method,
                     rpc_method.__rpc_method__.param_type,
                     rpc_method.__rpc_method__.cancelable,
+                    rpc_method.__rpc_method__.threaded,
                 )
                 for method, rpc_method in map(
                     lambda m1: (m1, cast(RpcMethod, m1)),
@@ -324,10 +324,11 @@ def add_method(
         func: Callable[..., Any],
         param_type: Optional[Type[Any]] = None,
         cancelable: bool = True,
+        threaded: bool = False,
     ) -> None:
         self.__ensure_initialized()
 
-        self.__methods[name] = RpcMethodEntry(name, func, param_type, cancelable)
+        self.__methods[name] = RpcMethodEntry(name, func, param_type, cancelable, threaded)
 
     def remove_method(self, name: str) -> Optional[RpcMethodEntry]:
         self.__ensure_initialized()
@@ -741,12 +742,10 @@ async def handle_request(self, message: JsonRPCRequest) -> None:
 
             params = self._convert_params(e.method, e.param_type, message.params)
 
-            is_threaded_method = is_threaded_callable(e.method)
-
-            if not is_threaded_method and not e.is_coroutine:
+            if not e.threaded and not e.is_coroutine:
                 self.send_response(message.id, e.method(*params[0], **params[1]))
             else:
-                if is_threaded_method:
+                if e.threaded:
                     if e.is_coroutine:
                         task = run_coroutine_in_thread(
                             ensure_coroutine(cast(Callable[..., Any], e.method)),
@@ -756,10 +755,7 @@ async def handle_request(self, message: JsonRPCRequest) -> None:
                     else:
                         task = asyncio.wrap_future(run_in_thread(e.method, *params[0], **params[1]))
                 else:
-                    task = create_sub_task(
-                        ensure_coroutine(e.method)(*params[0], **params[1]),
-                        name=message.method,
-                    )
+                    task = asyncio.create_task(e.method(*params[0], **params[1]), name=message.method)
 
                 with self._received_request_lock:
                     self._received_request[message.id] = ReceivedRequestEntry(task, message, e.cancelable)
@@ -845,20 +841,18 @@ async def handle_notification(self, message: JsonRPCNotification) -> None:
         try:
             params = self._convert_params(e.method, e.param_type, message.params)
 
-            if not e.is_coroutine:
+            if not e.threaded and not e.is_coroutine:
                 e.method(*params[0], **params[1])
             else:
-                if is_threaded_callable(e.method):
-                    task = run_coroutine_in_thread(
-                        ensure_coroutine(cast(Callable[..., Any], e.method)),
-                        *params[0],
-                        **params[1],
-                    )
+                if e.threaded:
+                    if e.is_coroutine:
+                        task = run_coroutine_in_thread(
+                            ensure_coroutine(cast(Callable[..., Any], e.method)), *params[0], **params[1]
+                        )
+                    else:
+                        task = asyncio.wrap_future(run_in_thread(e.method, *params[0], **params[1]))
                 else:
-                    task = create_sub_task(
-                        ensure_coroutine(e.method)(*params[0], **params[1]),
-                        name=message.method,
-                    )
+                    task = asyncio.create_task(e.method(*params[0], **params[1]), name=message.method)
 
                 await task
         except asyncio.CancelledError:

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/code_action.py b/packages/language_server/src/robotcode/language_server/common/parts/code_action.py
@@ -2,7 +2,7 @@
 from itertools import chain
 from typing import TYPE_CHECKING, Any, Final, List, Optional, Union, cast
 
-from robotcode.core.concurrent import check_current_thread_canceled, threaded
+from robotcode.core.concurrent import check_current_thread_canceled
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     CodeAction,
@@ -66,8 +66,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
                 resolve_provider=len(self.resolve) > 0,
             )
 
-    @rpc_method(name="textDocument/codeAction", param_type=CodeActionParams)
-    @threaded
+    @rpc_method(name="textDocument/codeAction", param_type=CodeActionParams, threaded=True)
     def _text_document_code_action(
         self,
         text_document: TextDocumentIdentifier,
@@ -109,8 +108,7 @@ def _text_document_code_action(
 
         return results
 
-    @rpc_method(name="codeAction/resolve", param_type=CodeAction)
-    @threaded
+    @rpc_method(name="codeAction/resolve", param_type=CodeAction, threaded=True)
     def _text_document_code_action_resolve(self, params: CodeAction, *args: Any, **kwargs: Any) -> CodeAction:
         results: List[CodeAction] = []
 

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/code_lens.py b/packages/language_server/src/robotcode/language_server/common/parts/code_lens.py
@@ -1,12 +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,
-    run_in_thread,
-    threaded,
-)
+from robotcode.core.concurrent import FutureEx, check_current_thread_canceled, run_in_thread
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     CodeLens,
@@ -46,8 +41,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if len(self.collect):
             capabilities.code_lens_provider = CodeLensOptions(resolve_provider=True if len(self.resolve) > 0 else None)
 
-    @rpc_method(name="textDocument/codeLens", param_type=CodeLensParams)
-    @threaded
+    @rpc_method(name="textDocument/codeLens", param_type=CodeLensParams, threaded=True)
     def _text_document_code_lens(
         self, text_document: TextDocumentIdentifier, *args: Any, **kwargs: Any
     ) -> Optional[List[CodeLens]]:
@@ -74,8 +68,7 @@ def _text_document_code_lens(
 
         return results
 
-    @rpc_method(name="codeLens/resolve", param_type=CodeLens)
-    @threaded
+    @rpc_method(name="codeLens/resolve", param_type=CodeLens, threaded=True)
     def _code_lens_resolve(self, params: CodeLens, *args: Any, **kwargs: Any) -> CodeLens:
         results: List[CodeLens] = []
 

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/commands.py b/packages/language_server/src/robotcode/language_server/common/parts/commands.py
@@ -13,7 +13,6 @@
     cast,
 )
 
-from robotcode.core.concurrent import threaded
 from robotcode.core.lsp.types import (
     ErrorCodes,
     ExecuteCommandOptions,
@@ -82,8 +81,7 @@ def get_command_name(self, callback: _FUNC_TYPE, name: Optional[str] = None) ->
     def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         capabilities.execute_command_provider = ExecuteCommandOptions(list(self.commands.keys()))
 
-    @rpc_method(name="workspace/executeCommand", param_type=ExecuteCommandParams)
-    @threaded
+    @rpc_method(name="workspace/executeCommand", param_type=ExecuteCommandParams, threaded=True)
     def _workspace_execute_command(
         self,
         command: str,

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/completion.py b/packages/language_server/src/robotcode/language_server/common/parts/completion.py
@@ -2,7 +2,7 @@
 from itertools import chain
 from typing import TYPE_CHECKING, Any, Final, List, Optional, Union, cast
 
-from robotcode.core.concurrent import check_current_thread_canceled, threaded
+from robotcode.core.concurrent import check_current_thread_canceled
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     CompletionContext,
@@ -84,8 +84,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
                 completion_item=CompletionOptionsCompletionItemType(label_details_support=True),
             )
 
-    @rpc_method(name="textDocument/completion", param_type=CompletionParams)
-    @threaded
+    @rpc_method(name="textDocument/completion", param_type=CompletionParams, threaded=True)
     def _text_document_completion(
         self,
         text_document: TextDocumentIdentifier,
@@ -159,8 +158,7 @@ def update_completion_item_to_utf16(self, document: TextDocument, item: Completi
             item.text_edit.insert = document.range_to_utf16(item.text_edit.insert)
             item.text_edit.replace = document.range_to_utf16(item.text_edit.replace)
 
-    @rpc_method(name="completionItem/resolve", param_type=CompletionItem)
-    @threaded
+    @rpc_method(name="completionItem/resolve", param_type=CompletionItem, threaded=True)
     def _completion_item_resolve(self, params: CompletionItem, *args: Any, **kwargs: Any) -> CompletionItem:
         results: List[CompletionItem] = []
 

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/declaration.py b/packages/language_server/src/robotcode/language_server/common/parts/declaration.py
@@ -1,7 +1,7 @@
 from asyncio import CancelledError
 from typing import TYPE_CHECKING, Any, Final, List, Optional, Union
 
-from robotcode.core.concurrent import check_current_thread_canceled, threaded
+from robotcode.core.concurrent import check_current_thread_canceled
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     DeclarationParams,
@@ -49,8 +49,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if len(self.collect):
             capabilities.declaration_provider = True
 
-    @rpc_method(name="textDocument/declaration", param_type=DeclarationParams)
-    @threaded
+    @rpc_method(name="textDocument/declaration", param_type=DeclarationParams, threaded=True)
     def _text_document_declaration(
         self,
         text_document: TextDocumentIdentifier,

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/definition.py b/packages/language_server/src/robotcode/language_server/common/parts/definition.py
@@ -1,7 +1,7 @@
 from concurrent.futures import CancelledError
 from typing import TYPE_CHECKING, Any, Final, List, Optional, Union
 
-from robotcode.core.concurrent import check_current_thread_canceled, threaded
+from robotcode.core.concurrent import check_current_thread_canceled
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     DefinitionParams,
@@ -49,8 +49,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if len(self.collect):
             capabilities.definition_provider = True
 
-    @rpc_method(name="textDocument/definition", param_type=DefinitionParams)
-    @threaded
+    @rpc_method(name="textDocument/definition", param_type=DefinitionParams, threaded=True)
     def _text_document_definition(
         self,
         text_document: TextDocumentIdentifier,

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/diagnostics.py b/packages/language_server/src/robotcode/language_server/common/parts/diagnostics.py
@@ -390,8 +390,7 @@ def publish_diagnostics(self, document: TextDocument, diagnostics: List[Diagnost
     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)
-    @threaded
+    @rpc_method(name="textDocument/diagnostic", param_type=DocumentDiagnosticParams, threaded=True)
     def _text_document_diagnostic(
         self,
         text_document: TextDocumentIdentifier,
@@ -431,8 +430,7 @@ def get_diagnostics_data(self, document: TextDocument) -> DiagnosticsData:
 
         return data
 
-    @rpc_method(name="workspace/diagnostic", param_type=WorkspaceDiagnosticParams)
-    @threaded
+    @rpc_method(name="workspace/diagnostic", param_type=WorkspaceDiagnosticParams, threaded=True)
     def _workspace_diagnostic(
         self,
         identifier: Optional[str],

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/document_highlight.py b/packages/language_server/src/robotcode/language_server/common/parts/document_highlight.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 check_current_thread_canceled
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     DocumentHighlight,
@@ -37,11 +37,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
     def collect(sender, document: TextDocument, position: Position) -> Optional[List[DocumentHighlight]]:  # NOSONAR
         ...
 
-    @rpc_method(
-        name="textDocument/documentHighlight",
-        param_type=DocumentHighlightParams,
-    )
-    @threaded
+    @rpc_method(name="textDocument/documentHighlight", param_type=DocumentHighlightParams, threaded=True)
     def _text_document_document_highlight(
         self,
         text_document: TextDocumentIdentifier,
@@ -61,6 +57,8 @@ def _text_document_document_highlight(
             document.position_from_utf16(position),
             callback_filter=language_id_filter(document),
         ):
+            check_current_thread_canceled()
+
             if isinstance(result, BaseException):
                 if not isinstance(result, CancelledError):
                     self._logger.exception(result, exc_info=result)

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/document_symbols.py b/packages/language_server/src/robotcode/language_server/common/parts/document_symbols.py
@@ -97,7 +97,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
                 else:
                     capabilities.document_symbol_provider = True
 
-    @rpc_method(name="textDocument/documentSymbol", param_type=DocumentSymbolParams)
+    @rpc_method(name="textDocument/documentSymbol", param_type=DocumentSymbolParams, threaded=True)
     @threaded
     def _text_document_symbol(
         self, text_document: TextDocumentIdentifier, *args: Any, **kwargs: Any

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/folding_range.py b/packages/language_server/src/robotcode/language_server/common/parts/folding_range.py
@@ -1,7 +1,6 @@
 from concurrent.futures import CancelledError
 from typing import TYPE_CHECKING, Any, Final, List, Optional
 
-from robotcode.core.concurrent import threaded
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     FoldingRange,
@@ -38,8 +37,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if len(self.collect):
             capabilities.folding_range_provider = True
 
-    @rpc_method(name="textDocument/foldingRange", param_type=FoldingRangeParams)
-    @threaded
+    @rpc_method(name="textDocument/foldingRange", param_type=FoldingRangeParams, threaded=True)
     def _text_document_folding_range(
         self, text_document: TextDocumentIdentifier, *args: Any, **kwargs: Any
     ) -> Optional[List[FoldingRange]]:

diff --git a/packages/language_server/src/robotcode/language_server/common/parts/formatting.py b/packages/language_server/src/robotcode/language_server/common/parts/formatting.py
@@ -1,7 +1,7 @@
 from concurrent.futures 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 check_current_thread_canceled
 from robotcode.core.event import event
 from robotcode.core.lsp.types import (
     DocumentFormattingOptions,
@@ -58,8 +58,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if len(self.format_range):
             capabilities.document_range_formatting_provider = DocumentRangeFormattingOptions(work_done_progress=True)
 
-    @rpc_method(name="textDocument/formatting", param_type=DocumentFormattingParams)
-    @threaded
+    @rpc_method(name="textDocument/formatting", param_type=DocumentFormattingParams, threaded=True)
     def _text_document_formatting(
         self,
         params: DocumentFormattingParams,
@@ -96,11 +95,7 @@ def _text_document_formatting(
 
         return None
 
-    @rpc_method(
-        name="textDocument/rangeFormatting",
-        param_type=DocumentRangeFormattingParams,
-    )
-    @threaded
+    @rpc_method(name="textDocument/rangeFormatting", param_type=DocumentRangeFormattingParams, threaded=True)
     def _text_document_range_formatting(
         self,
         params: DocumentFormattingParams,