reimplement run locking and starting coroutines in another thread

Author: d-biehl

robotcode/jsonrpc2/protocol.py

@@ -595,7 +595,12 @@ def handle_response(self, message: JsonRPCResponse) -> None:
                     entry.future.set_result(res)
                 else:
                     if entry.future._loop.is_running():
-                        entry.future._loop.call_soon_threadsafe(entry.future.set_result, res)
+
+                        def s(f: asyncio.Future[Any], r: Any) -> None:
+                            if not f.done():
+                                f.set_result(r)
+
+                        entry.future._loop.call_soon_threadsafe(s, entry.future, res)
                     else:
                         self._logger.warning("Response loop is not running.")
 
@@ -700,8 +705,8 @@ def done(t: asyncio.Task[Any]) -> None:
                 self.send_error(JsonRPCErrors.REQUEST_CANCELLED, "Request canceled.", id=message.id)
             except (SystemExit, KeyboardInterrupt):
                 raise
-            except JsonRPCErrorException as ex:
-                self.send_error(ex.code, ex.message, id=message.id, data=ex.data)
+            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)

robotcode/language_server/common/parts/diagnostics.py

@@ -90,6 +90,8 @@ def __init__(self, protocol: LanguageServerProtocol) -> None:
         ] = {}
 
         self._current_workspace_task: Optional[asyncio.Task[WorkspaceDiagnosticReport]] = None
+        self.in_get_document_diagnostics = Event(True)
+        self.in_get_workspace_diagnostics = Event(True)
 
     def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if (
@@ -108,14 +110,26 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
     async def collect(sender, document: TextDocument) -> DiagnosticsResult:  # NOSONAR
         ...
 
+    @async_tasking_event_iterator
+    async def collect_stage2(sender, document: TextDocument) -> DiagnosticsResult:  # NOSONAR
+        ...
+
     @async_tasking_event
-    async def collect_workspace_documents(sender) -> List[WorkspaceDocumentsResult]:  # NOSONAR
+    async def load_workspace_documents(sender) -> List[WorkspaceDocumentsResult]:  # NOSONAR
         ...
 
     @async_tasking_event
     async def on_workspace_loaded(sender) -> None:  # NOSONAR
         ...
 
+    @async_tasking_event
+    async def on_workspace_diagnostics_ended(sender) -> None:  # NOSONAR
+        ...
+
+    @async_tasking_event
+    async def on_document_diagnostics_ended(sender) -> None:  # NOSONAR
+        ...
+
     @async_event
     async def on_get_analysis_progress_mode(sender, uri: Uri) -> Optional[AnalysisProgressMode]:  # NOSONAR
         ...
@@ -125,17 +139,18 @@ async def on_get_diagnostics_mode(sender, uri: Uri) -> Optional[DiagnosticsMode]
         ...
 
     async def ensure_workspace_loaded(self) -> None:
-        async with self._workspace_load_lock:
-            if not self._workspace_loaded:
-                if self.workspace_loaded_event.is_set():
-                    return
-
-                try:
-                    await self.collect_workspace_documents(self)
-                finally:
-                    self._workspace_loaded = True
-                    self.workspace_loaded_event.set()
-                    await self.on_workspace_loaded(self)
+        if not self._workspace_loaded:
+            async with self._workspace_load_lock:
+                if not self._workspace_loaded:
+                    if self.workspace_loaded_event.is_set():
+                        return
+
+                    try:
+                        await self.load_workspace_documents(self)
+                    finally:
+                        self._workspace_loaded = True
+                        self.workspace_loaded_event.set()
+                        await self.on_workspace_loaded(self)
 
     async def get_document_diagnostics(self, document: TextDocument) -> RelatedFullDocumentDiagnosticReport:
         return await document.get_cache(self.__get_document_diagnostics)
@@ -171,11 +186,13 @@ async def _text_document_diagnostic(
         *args: Any,
         **kwargs: Any,
     ) -> DocumentDiagnosticReport:
-        self._logger.info(lambda: f"textDocument/diagnostic for {text_document}")
+        self._logger.debug(lambda: f"textDocument/diagnostic for {text_document}")
 
-        await self.ensure_workspace_loaded()
+        self.in_get_document_diagnostics.clear()
 
         try:
+            await self.ensure_workspace_loaded()
+
             document = await self.parent.documents.get(text_document.uri)
             if document is None:
                 raise JsonRPCErrorException(ErrorCodes.INVALID_PARAMS, f"Document {text_document!r} not found")
@@ -187,7 +204,7 @@ async def _get_diagnostics() -> Optional[RelatedFullDocumentDiagnosticReport]:
                 return None
 
             if document in self._current_document_tasks and not self._current_document_tasks[document].done():
-                self._logger.info(lambda: f"textDocument/diagnostic cancel old task {text_document}")
+                self._logger.debug(lambda: f"textDocument/diagnostic cancel old task {text_document}")
                 self._current_document_tasks[document].cancel()
 
             task = create_sub_task(_get_diagnostics())
@@ -199,7 +216,7 @@ async def _get_diagnostics() -> Optional[RelatedFullDocumentDiagnosticReport]:
                     raise RuntimeError("Unexpected result.")
 
             except asyncio.CancelledError as e:
-                self._logger.info(lambda: f"textDocument/diagnostic canceled {text_document}")
+                self._logger.debug(lambda: f"textDocument/diagnostic canceled {text_document}")
 
                 raise JsonRPCErrorException(
                     ErrorCodes.SERVER_CANCELLED, "Cancelled", data=DiagnosticServerCancellationData(True)
@@ -213,7 +230,10 @@ async def _get_diagnostics() -> Optional[RelatedFullDocumentDiagnosticReport]:
 
             return result
         finally:
-            self._logger.info(lambda: f"textDocument/diagnostic ready  {text_document}")
+            self.in_get_document_diagnostics.set()
+            await self.on_document_diagnostics_ended(self)
+
+            self._logger.debug(lambda: f"textDocument/diagnostic ready  {text_document}")
 
     @rpc_method(name="workspace/diagnostic", param_type=WorkspaceDiagnosticParams)
     @threaded()
@@ -226,9 +246,7 @@ async def _workspace_diagnostic(
         *args: Any,
         **kwargs: Any,
     ) -> WorkspaceDiagnosticReport:
-        self._logger.critical("workspace/diagnostic")
-
-        await self.ensure_workspace_loaded()
+        self._logger.debug("workspace/diagnostic")
 
         async def _get_diagnostics() -> WorkspaceDiagnosticReport:
             result: List[WorkspaceDocumentDiagnosticReport] = []
@@ -303,22 +321,30 @@ async def _task(doc: TextDocument) -> None:
 
                 return WorkspaceDiagnosticReport(items=[])
 
-        if self._current_workspace_task is not None:
-            self._current_workspace_task.cancel()
-
-        task = create_sub_task(_get_diagnostics() if partial_result_token is None else _get_partial_diagnostics())
-        self._current_workspace_task = task
+        await self.in_get_document_diagnostics.wait()
+        self.in_get_workspace_diagnostics.clear()
         try:
-            return await task
-        except asyncio.CancelledError as e:
-            self._logger.critical("workspace/diagnostic canceled")
-            raise JsonRPCErrorException(
-                ErrorCodes.SERVER_CANCELLED, "Cancelled", data=DiagnosticServerCancellationData(True)
-            ) from e
+            await self.ensure_workspace_loaded()
+
+            if self._current_workspace_task is not None:
+                self._current_workspace_task.cancel()
+
+            task = create_sub_task(_get_diagnostics() if partial_result_token is None else _get_partial_diagnostics())
+            self._current_workspace_task = task
+            try:
+                return await task
+            except asyncio.CancelledError as e:
+                self._logger.debug("workspace/diagnostic canceled")
+                raise JsonRPCErrorException(
+                    ErrorCodes.SERVER_CANCELLED, "ServerCancelled", data=DiagnosticServerCancellationData(True)
+                ) from e
+            finally:
+                if self._current_workspace_task == task:
+                    self._current_workspace_task = None
+                self._logger.debug("workspace/diagnostic ready")
         finally:
-            if self._current_workspace_task == task:
-                self._current_workspace_task = None
-            self._logger.critical("workspace/diagnostic ready")
+            self.in_get_workspace_diagnostics.set()
+            await self.on_workspace_diagnostics_ended(self)
 
     def cancel_workspace_diagnostics(self) -> None:
         if self._current_workspace_task is not None and not self._current_workspace_task.done():

robotcode/language_server/robotframework/parts/codelens.py

@@ -3,7 +3,7 @@
 import ast
 from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple, cast
 
-from ....utils.async_tools import create_sub_task, run_coroutine_in_thread, threaded
+from ....utils.async_tools import create_sub_task, threaded
 from ....utils.logging import LoggingDescriptor
 from ...common.decorators import language_id
 from ...common.lsp_types import CodeLens, Command
@@ -28,16 +28,13 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         parent.code_lens.collect.add(self.collect)
         parent.code_lens.resolve.add(self.resolve)
 
-        parent.robot_references.cache_cleared.add(self.robot_references_cache_cleared)
-
         self._running_task: Set[Tuple[TextDocument, KeywordDoc]] = set()
 
-        parent.diagnostics.on_workspace_loaded.add(self.diagnostics_on_workspace_loaded)
-
-    async def robot_references_cache_cleared(self, sender: Any) -> None:  # NOSONAR
-        await self.parent.code_lens.refresh()
+        parent.diagnostics.on_workspace_loaded.add(self.codelens_refresh)
+        parent.diagnostics.on_document_diagnostics_ended.add(self.codelens_refresh)
+        parent.diagnostics.on_workspace_diagnostics_ended.add(self.codelens_refresh)
 
-    async def diagnostics_on_workspace_loaded(self, sender: Any) -> None:  # NOSONAR
+    async def codelens_refresh(self, sender: Any) -> None:  # NOSONAR
         await self.parent.code_lens.refresh()
 
     @language_id("robotframework")
@@ -114,7 +111,11 @@ async def resolve(self, sender: Any, code_lens: CodeLens) -> Optional[CodeLens]:
         name = code_lens.data["name"]
         line = code_lens.data["line"]
 
-        if self.parent.diagnostics.workspace_loaded_event.is_set():
+        if (
+            self.parent.diagnostics.workspace_loaded_event.is_set()
+            and self.parent.diagnostics.in_get_workspace_diagnostics.is_set()
+            and self.parent.diagnostics.in_get_document_diagnostics.is_set()
+        ):
             kw_doc = await self.get_keyword_definition_at_line(namespace, name, line)
 
             if kw_doc is not None and not kw_doc.is_error_handler:
@@ -131,17 +132,17 @@ async def find_refs() -> None:
                         if document is None or kw_doc is None:
                             return
 
-                        await run_coroutine_in_thread(
-                            self.parent.robot_references.find_keyword_references,
-                            document,
-                            kw_doc,
-                            include_declaration=False,
-                        )
-
-                        # await self.parent.robot_references.find_keyword_references(
-                        #     document, kw_doc, include_declaration=False
+                        # await run_coroutine_in_thread(
+                        #     self.parent.robot_references.find_keyword_references,
+                        #     document,
+                        #     kw_doc,
+                        #     include_declaration=False,
                         # )
 
+                        await self.parent.robot_references.find_keyword_references(
+                            document, kw_doc, include_declaration=False
+                        )
+
                         await self.parent.code_lens.refresh()
 
                     key = (document, kw_doc)

robotcode/language_server/robotframework/parts/diagnostics.py

@@ -43,19 +43,17 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
 
         parent.diagnostics.collect.add(self.collect_namespace_diagnostics)
 
-        parent.diagnostics.collect.add(self.collect_unused_references)
+        parent.diagnostics.collect_stage2.add(self.collect_unused_references)
 
         parent.documents_cache.namespace_invalidated.add(self.namespace_invalidated)
 
-    async def cache_cleared(self, sender: Any) -> None:
-        ...
-
     @language_id("robotframework")
     async def namespace_invalidated(self, sender: Any, document: TextDocument) -> None:
         self.parent.diagnostics.cancel_workspace_diagnostics()
 
     @language_id("robotframework")
     @threaded()
+    @_logger.call
     async def collect_namespace_diagnostics(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
         try:
             namespace = await self.parent.documents_cache.get_namespace(document)
@@ -120,6 +118,7 @@ def _create_error_from_token(self, token: Token, source: Optional[str] = None) -
         )
 
     @language_id("robotframework")
+    @threaded()
     @_logger.call
     async def collect_token_errors(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
         from robot.errors import VariableError

robotcode/language_server/robotframework/parts/robot_workspace.py

@@ -32,7 +32,7 @@ class RobotWorkspaceProtocolPart(RobotLanguageServerProtocolPart):
     def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         super().__init__(parent)
         self.parent.documents.on_read_document_text.add(self._on_read_document_text)
-        self.parent.diagnostics.collect_workspace_documents.add(self._collect_workspace_documents)
+        self.parent.diagnostics.load_workspace_documents.add(self._load_workspace_documents)
         self.parent.diagnostics.on_get_diagnostics_mode.add(self.on_get_diagnostics_mode)
         self.parent.diagnostics.on_get_analysis_progress_mode.add(self.on_get_analysis_progress_mode)
         self.workspace_loaded = Event()
@@ -53,7 +53,7 @@ async def on_get_analysis_progress_mode(self, sender: Any, uri: Uri) -> Optional
         return config.progress_mode
 
     @threaded()
-    async def _collect_workspace_documents(self, sender: Any) -> List[WorkspaceDocumentsResult]:
+    async def _load_workspace_documents(self, sender: Any) -> List[WorkspaceDocumentsResult]:
 
         result: List[WorkspaceDocumentsResult] = []
 

robotcode/utils/async_cache.py

@@ -27,30 +27,33 @@ def __init__(self, max_items: Optional[int] = 128) -> None:
         self._order: Optional[List[Tuple[Any, ...]]] = None
         if self.max_items:
             self._order = []
+        self.lock: Lock = Lock()
 
     async def has(self, *args: Any, **kwargs: Any) -> bool:
         key = self._make_key(*args, **kwargs)
-
-        if key in self._cache:
-            return self._cache[key].has_data
+        async with self.lock:
+            if key in self._cache:
+                return self._cache[key].has_data
 
         return False
 
     async def get(self, func: Callable[..., Awaitable[_T]], *args: Any, **kwargs: Any) -> _T:
         key = self._make_key(*args, **kwargs)
 
-        entry = self._cache[key]
+        async with self.lock:
+            entry = self._cache[key]
 
         async with entry.lock:
             if not entry.has_data:
                 entry.data = await func(*args, **kwargs)
                 entry.has_data = True
 
-                if self._order is not None and self.max_items is not None:
-                    self._order.insert(0, key)
+                async with self.lock:
+                    if self._order is not None and self.max_items is not None:
+                        self._order.insert(0, key)
 
-                    if len(self._order) > self.max_items:
-                        del self._cache[self._order.pop()]
+                        if len(self._order) > self.max_items:
+                            del self._cache[self._order.pop()]
 
             return cast(_T, entry.data)
 
@@ -59,6 +62,7 @@ def _make_key(*args: Any, **kwargs: Any) -> Tuple[Any, ...]:
         return (tuple(_freeze(v) for v in args), hash(frozenset({k: _freeze(v) for k, v in kwargs.items()})))
 
     async def clear(self) -> None:
-        self._cache.clear()
-        if self._order is not None:
-            self._order.clear()
+        async with self.lock:
+            self._cache.clear()
+            if self._order is not None:
+                self._order.clear()