refactor(langserver): move core.utils.threading to core.concurrent and rename some function

Author: d-biehl

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

@@ -31,8 +31,8 @@
     cast,
 )
 
+from .concurrent import is_threaded_callable
 from .utils.inspect import ensure_coroutine
-from .utils.threading import is_threaded_callable
 
 _T = TypeVar("_T")
 

packages/core/src/robotcode/core/utils/threading.py renamed to packages/core/src/robotcode/core/concurrent.py

@@ -1,26 +1,33 @@
 import inspect
 from concurrent.futures import CancelledError, Future
-from threading import Event, Thread, current_thread, local
-from typing import Any, Callable, Dict, Generic, Optional, Tuple, TypeVar, cast, overload
+from threading import Event, RLock, Thread, current_thread, local
+from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, TypeVar, cast, overload
 
 _F = TypeVar("_F", bound=Callable[..., Any])
 _TResult = TypeVar("_TResult")
 
-__THREADED_MARKER = "__threaded__"
+__THREADED_MARKER = "__robotcode_threaded"
 
 
 class FutureEx(Future, Generic[_TResult]):  # type: ignore[type-arg]
     def __init__(self) -> None:
         super().__init__()
-        self.cancelation_requested = Event()
+        self.cancelation_requested_event = Event()
+
+    @property
+    def cancelation_requested(self) -> bool:
+        return self.cancelation_requested_event.is_set()
 
     def cancel(self) -> bool:
-        self.cancelation_requested.set()
+        self.cancelation_requested_event.set()
         return super().cancel()
 
     def result(self, timeout: Optional[float] = None) -> _TResult:
         return cast(_TResult, super().result(timeout))
 
+    def add_done_callback(self, fn: Callable[["FutureEx[Any]"], Any]) -> None:
+        super().add_done_callback(fn)  # type: ignore[arg-type]
+
 
 @overload
 def threaded(__func: _F) -> _F:
@@ -63,44 +70,72 @@ def __init__(self) -> None:
 def _run_callable_in_thread_handler(
     future: FutureEx[_TResult], callable: Callable[..., _TResult], args: Tuple[Any, ...], kwargs: Dict[str, Any]
 ) -> None:
+    if not future.set_running_or_notify_cancel():
+        return
+
     _local_storage._local_future = future
-    future.set_running_or_notify_cancel()
+
     try:
         result = callable(*args, **kwargs)
     except Exception as e:
         # TODO: add traceback to exception e.traceback = format_exc()
-        if not future.cancelled():
-            future.set_exception(e)
+        future.set_exception(e)
     else:
-        if not future.cancelled():
-            future.set_result(result)
+        future.set_result(result)
     finally:
         _local_storage._local_future = None
 
 
-def is_thread_cancelled() -> bool:
+def is_current_thread_cancelled() -> bool:
     local_future = _local_storage._local_future
-    return local_future is not None and local_future.cancelation_requested.is_set()
+    return local_future is not None and local_future.cancelation_requested
 
 
-def check_thread_canceled(at_least_seconds: Optional[float] = None) -> None:
+def check_current_thread_canceled(at_least_seconds: Optional[float] = None) -> None:
     local_future = _local_storage._local_future
     if local_future is None:
         return
 
     if at_least_seconds is None or at_least_seconds <= 0:
-        if not local_future.cancelation_requested.is_set():
+        if not local_future.cancelation_requested:
             return
-    elif not local_future.cancelation_requested.wait(at_least_seconds):
+    elif not local_future.cancelation_requested_event.wait(at_least_seconds):
         return
 
     name = current_thread().name
     raise CancelledError(f"Thread {name+' ' if name else ' '}cancelled")
 
 
-def run_callable_in_thread(callable: Callable[..., _TResult], *args: Any, **kwargs: Any) -> FutureEx[_TResult]:
+_running_callables_lock = RLock()
+_running_callables: Dict[FutureEx[Any], Thread] = {}
+
+
+def _remove_future_from_running_callables(future: FutureEx[Any]) -> None:
+    with _running_callables_lock:
+        _running_callables.pop(future, None)
+
+
+def run_in_thread(callable: Callable[..., _TResult], *args: Any, **kwargs: Any) -> FutureEx[_TResult]:
     future: FutureEx[_TResult] = FutureEx()
+    with _running_callables_lock:
+        thread = Thread(
+            target=_run_callable_in_thread_handler, args=(future, callable, args, kwargs), name=str(callable)
+        )
+        _running_callables[future] = thread
+        future.add_done_callback(_remove_future_from_running_callables)
 
-    Thread(target=_run_callable_in_thread_handler, args=(future, callable, args, kwargs), name=str(callable)).start()
+    thread.start()
 
     return future
+
+
+def cancel_running_callables(timeout: Optional[float] = None) -> None:
+    threads: List[Thread] = []
+    with _running_callables_lock:
+        for future, thread in _running_callables.items():
+            if not future.cancelation_requested:
+                future.cancel()
+                threads.append(thread)
+    for thread in threads:
+        if thread is not current_thread():
+            thread.join(timeout=timeout)

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

@@ -38,10 +38,10 @@
     create_sub_task,
     run_coroutine_in_thread,
 )
+from robotcode.core.concurrent import is_threaded_callable, run_in_thread
 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
-from robotcode.core.utils.threading import is_threaded_callable, run_callable_in_thread
 
 __all__ = [
     "JsonRPCErrors",
@@ -742,7 +742,7 @@ async def handle_request(self, message: JsonRPCRequest) -> None:
                             ensure_coroutine(cast(Callable[..., Any], e.method)), *params[0], **params[1]
                         )
                     else:
-                        task = asyncio.wrap_future(run_callable_in_thread(e.method, *params[0], **params[1]))
+                        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]),

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

@@ -2,6 +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.event import event
 from robotcode.core.lsp.types import (
     CodeAction,
@@ -14,7 +15,6 @@
     TextDocumentIdentifier,
 )
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import check_thread_canceled, threaded
 from robotcode.jsonrpc2.protocol import rpc_method
 from robotcode.language_server.common.decorators import CODE_ACTION_KINDS_ATTR, HasCodeActionKinds, language_id_filter
 from robotcode.language_server.common.parts.protocol_part import LanguageServerProtocolPart
@@ -82,7 +82,7 @@ def _text_document_code_action(
         for result in self.collect(
             self, document, document.range_from_utf16(range), context, callback_filter=language_id_filter(document)
         ):
-            check_thread_canceled()
+            check_current_thread_canceled()
 
             if isinstance(result, BaseException):
                 if not isinstance(result, CancelledError):
@@ -107,7 +107,7 @@ def _text_document_code_action_resolve(
         results: List[CodeAction] = []
 
         for result in self.resolve(self, params):
-            check_thread_canceled()
+            check_current_thread_canceled()
 
             if isinstance(result, BaseException):
                 if not isinstance(result, CancelledError):

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

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Any, Final, List, Optional
 
 from robotcode.core.async_tools import async_tasking_event, create_sub_task
+from robotcode.core.concurrent import threaded
 from robotcode.core.lsp.types import (
     CodeLens,
     CodeLensOptions,
@@ -13,7 +14,6 @@
     TextDocumentIdentifier,
 )
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import threaded
 from robotcode.jsonrpc2.protocol import rpc_method
 from robotcode.language_server.common.decorators import language_id_filter
 from robotcode.language_server.common.text_document import TextDocument

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

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Final, List, Optional, cast
 
+from robotcode.core.concurrent import threaded
 from robotcode.core.lsp.types import (
     ErrorCodes,
     ExecuteCommandOptions,
@@ -15,7 +16,6 @@
 )
 from robotcode.core.utils.dataclasses import from_dict
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import threaded
 from robotcode.jsonrpc2.protocol import JsonRPCErrorException, rpc_method
 from robotcode.language_server.common.decorators import get_command_id, is_command
 from robotcode.language_server.common.parts.protocol_part import LanguageServerProtocolPart

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

@@ -2,6 +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.event import event
 from robotcode.core.lsp.types import (
     CompletionContext,
@@ -18,7 +19,6 @@
     TextEdit,
 )
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import check_thread_canceled, threaded
 from robotcode.jsonrpc2.protocol import rpc_method
 from robotcode.language_server.common.decorators import (
     ALL_COMMIT_CHARACTERS_ATTR,
@@ -92,7 +92,7 @@ def _text_document_completion(
         results: List[Union[List[CompletionItem], CompletionList]] = []
 
         if context is not None and context.trigger_kind == CompletionTriggerKind.TRIGGER_CHARACTER:
-            check_thread_canceled(0.25)
+            check_current_thread_canceled(0.25)
 
         document = self.parent.documents.get(text_document.uri)
         if document is None:
@@ -101,7 +101,7 @@ def _text_document_completion(
         p = document.position_from_utf16(position)
 
         for result in self.collect(self, document, p, context, callback_filter=language_id_filter(document)):
-            check_thread_canceled()
+            check_current_thread_canceled()
 
             if isinstance(result, BaseException):
                 if not isinstance(result, CancelledError):
@@ -114,7 +114,7 @@ def _text_document_completion(
             return None
 
         for result in results:
-            check_thread_canceled()
+            check_current_thread_canceled()
 
             if isinstance(result, CompletionList):
                 for item in result.items:

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

@@ -3,6 +3,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.event import event
 from robotcode.core.lsp.types import (
     DeclarationParams,
@@ -13,7 +14,6 @@
     TextDocumentIdentifier,
 )
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import check_thread_canceled, threaded
 from robotcode.jsonrpc2.protocol import rpc_method
 from robotcode.language_server.common.decorators import language_id_filter
 from robotcode.language_server.common.parts.protocol_part import LanguageServerProtocolPart
@@ -64,7 +64,7 @@ def _text_document_declaration(
             return None
 
         for result in self.collect(self, document, position, callback_filter=language_id_filter(document)):
-            check_thread_canceled()
+            check_current_thread_canceled()
 
             if isinstance(result, BaseException):
                 if not isinstance(result, CancelledError):

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

@@ -3,6 +3,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.event import event
 from robotcode.core.lsp.types import (
     DefinitionParams,
@@ -13,7 +14,6 @@
     TextDocumentIdentifier,
 )
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import check_thread_canceled, threaded
 from robotcode.jsonrpc2.protocol import rpc_method
 from robotcode.language_server.common.decorators import language_id_filter
 from robotcode.language_server.common.parts.protocol_part import LanguageServerProtocolPart
@@ -62,7 +62,7 @@ def _text_document_definition(
         for result in self.collect(
             self, document, document.position_from_utf16(position), callback_filter=language_id_filter(document)
         ):
-            check_thread_canceled()
+            check_current_thread_canceled()
 
             if isinstance(result, BaseException):
                 if not isinstance(result, CancelledError):

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

@@ -17,6 +17,7 @@
     check_canceled,
     create_sub_task,
 )
+from robotcode.core.concurrent import threaded
 from robotcode.core.lsp.types import (
     Diagnostic,
     DiagnosticOptions,
@@ -35,7 +36,6 @@
 )
 from robotcode.core.uri import Uri
 from robotcode.core.utils.logging import LoggingDescriptor
-from robotcode.core.utils.threading import threaded
 from robotcode.jsonrpc2.protocol import JsonRPCErrorException, rpc_method
 from robotcode.language_server.common.decorators import language_id_filter
 from robotcode.language_server.common.parts.protocol_part import LanguageServerProtocolPart