optimize workspace diagnostics

Author: d-biehl

robotcode/language_server/common/parts/diagnostics.py

@@ -14,7 +14,6 @@
     async_tasking_event,
     async_tasking_event_iterator,
     create_sub_task,
-    threaded,
 )
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
@@ -92,6 +91,7 @@ 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)
+        self._collect_full_diagnostics = False
 
     def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if (
@@ -106,12 +106,17 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
                 work_done_progress=True,
             )
 
-    @async_tasking_event_iterator
-    async def collect(sender, document: TextDocument) -> DiagnosticsResult:  # NOSONAR
-        ...
+    @property
+    def collect_full_diagnostics(self) -> bool:
+        return self._collect_full_diagnostics
+
+    async def set_collect_full_diagnostics(self, value: bool) -> None:
+        if self._collect_full_diagnostics != value:
+            self._collect_full_diagnostics = value
+            await self.refresh()
 
     @async_tasking_event_iterator
-    async def collect_stage2(sender, document: TextDocument) -> DiagnosticsResult:  # NOSONAR
+    async def collect(sender, document: TextDocument, full: bool) -> DiagnosticsResult:  # NOSONAR
         ...
 
     @async_tasking_event
@@ -153,6 +158,8 @@ async def ensure_workspace_loaded(self) -> None:
                         await self.on_workspace_loaded(self)
 
     async def get_document_diagnostics(self, document: TextDocument) -> RelatedFullDocumentDiagnosticReport:
+        if self.collect_full_diagnostics:
+            return await document.get_cache(self.__get_full_document_diagnostics)
         return await document.get_cache(self.__get_document_diagnostics)
 
     async def __get_document_diagnostics(self, document: TextDocument) -> RelatedFullDocumentDiagnosticReport:
@@ -163,6 +170,29 @@ async def __get_document_diagnostics(self, document: TextDocument) -> RelatedFul
         async for result_any in self.collect(
             self,
             document,
+            full=False,
+            callback_filter=language_id_filter(document),
+            return_exceptions=True,
+        ):
+            result = cast(DiagnosticsResult, result_any)
+
+            if isinstance(result, BaseException):
+                if not isinstance(result, asyncio.CancelledError):
+                    self._logger.exception(result, exc_info=result)
+            else:
+                diagnostics.extend(result.diagnostics or [])
+
+        return RelatedFullDocumentDiagnosticReport(items=diagnostics, result_id=str(uuid.uuid4()))
+
+    async def __get_full_document_diagnostics(self, document: TextDocument) -> RelatedFullDocumentDiagnosticReport:
+        await self.ensure_workspace_loaded()
+
+        diagnostics: List[Diagnostic] = []
+
+        async for result_any in self.collect(
+            self,
+            document,
+            full=True,
             callback_filter=language_id_filter(document),
             return_exceptions=True,
         ):
@@ -177,7 +207,6 @@ async def __get_document_diagnostics(self, document: TextDocument) -> RelatedFul
         return RelatedFullDocumentDiagnosticReport(items=diagnostics, result_id=str(uuid.uuid4()))
 
     @rpc_method(name="textDocument/diagnostic", param_type=DocumentDiagnosticParams)
-    @threaded()
     async def _text_document_diagnostic(
         self,
         text_document: TextDocumentIdentifier,
@@ -236,7 +265,6 @@ async def _get_diagnostics() -> Optional[RelatedFullDocumentDiagnosticReport]:
             self._logger.debug(lambda: f"textDocument/diagnostic ready  {text_document}")
 
     @rpc_method(name="workspace/diagnostic", param_type=WorkspaceDiagnosticParams)
-    @threaded()
     async def _workspace_diagnostic(
         self,
         identifier: Optional[str],
@@ -272,7 +300,9 @@ async def _get_diagnostics() -> WorkspaceDiagnosticReport:
 
         async def _get_partial_diagnostics() -> WorkspaceDiagnosticReport:
             async with self.parent.window.progress(
-                "Analyse Workspace", progress_token=work_done_token, cancellable=False
+                f"Analyse {'full ' if self.collect_full_diagnostics else ''} Workspace",
+                progress_token=work_done_token,
+                cancellable=False,
             ) as progress:
 
                 async def _task(doc: TextDocument) -> None:
@@ -285,7 +315,7 @@ async def _task(doc: TextDocument) -> None:
                         else:
                             name = path.relative_to(folder.uri.to_path())
 
-                        progress.report(f"Analyse {name}")
+                        progress.report(f"Analyse {'full ' if self.collect_full_diagnostics else ''} {name}")
 
                     doc_result = await self.get_document_diagnostics(doc)
 
@@ -332,7 +362,9 @@ async def _task(doc: TextDocument) -> None:
             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
+                result = await task
+                await self.set_collect_full_diagnostics(True)
+                return result
             except asyncio.CancelledError as e:
                 self._logger.debug("workspace/diagnostic canceled")
                 raise JsonRPCErrorException(
@@ -350,6 +382,12 @@ def cancel_workspace_diagnostics(self) -> None:
         if self._current_workspace_task is not None and not self._current_workspace_task.done():
             self._current_workspace_task.cancel()
 
+    def cancel_document_diagnostics(self, document: TextDocument) -> None:
+        task = self._current_document_tasks.get(document, None)
+        if task is not None:
+            self._logger.critical(lambda: f"textDocument/diagnostic canceled {document}")
+            task.cancel()
+
     async def get_analysis_progress_mode(self, uri: Uri) -> AnalysisProgressMode:
         for e in await self.on_get_analysis_progress_mode(self, uri):
             if e is not None:

robotcode/language_server/robotframework/parts/diagnostics.py

@@ -43,18 +43,23 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
 
         parent.diagnostics.collect.add(self.collect_namespace_diagnostics)
 
-        parent.diagnostics.collect_stage2.add(self.collect_unused_references)
+        parent.diagnostics.collect.add(self.collect_unused_references)
 
         parent.documents_cache.namespace_invalidated.add(self.namespace_invalidated)
 
     @language_id("robotframework")
     async def namespace_invalidated(self, sender: Any, document: TextDocument) -> None:
         self.parent.diagnostics.cancel_workspace_diagnostics()
+        self.parent.diagnostics.cancel_document_diagnostics(document)
+        await self.parent.diagnostics.set_collect_full_diagnostics(False)
 
     @language_id("robotframework")
     @threaded()
     @_logger.call
-    async def collect_namespace_diagnostics(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_namespace_diagnostics(self, sender: Any, document: TextDocument, full: bool) -> DiagnosticsResult:
+        return await document.get_cache(self._collect_namespace_diagnostics, full)
+
+    async def _collect_namespace_diagnostics(self, document: TextDocument, full: bool) -> DiagnosticsResult:
         try:
             namespace = await self.parent.documents_cache.get_namespace(document)
             if namespace is None:
@@ -120,7 +125,10 @@ 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:
+    async def collect_token_errors(self, sender: Any, document: TextDocument, full: bool) -> DiagnosticsResult:
+        return await document.get_cache(self._collect_token_errors, full)
+
+    async def _collect_token_errors(self, document: TextDocument, full: bool) -> DiagnosticsResult:
         from robot.errors import VariableError
         from robot.parsing.lexer.tokens import Token
 
@@ -186,7 +194,10 @@ async def collect_token_errors(self, sender: Any, document: TextDocument) -> Dia
     @language_id("robotframework")
     @threaded()
     @_logger.call
-    async def collect_model_errors(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_model_errors(self, sender: Any, document: TextDocument, full: bool) -> DiagnosticsResult:
+        return await document.get_cache(self._collect_model_errors, full)
+
+    async def _collect_model_errors(self, document: TextDocument, full: bool) -> DiagnosticsResult:
 
         from ..utils.ast_utils import HasError, HasErrors
         from ..utils.async_ast import iter_nodes
@@ -235,7 +246,15 @@ async def collect_model_errors(self, sender: Any, document: TextDocument) -> Dia
     @language_id("robotframework")
     @threaded()
     @_logger.call
-    async def collect_unused_references(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_unused_references(self, sender: Any, document: TextDocument, full: bool) -> DiagnosticsResult:
+        if not full:
+            return DiagnosticsResult(self.collect_unused_references, [])
+        return await document.get_cache(self._collect_unused_references, full)
+
+    async def _collect_unused_references(self, document: TextDocument, full: bool) -> DiagnosticsResult:
+        if not full:
+            return DiagnosticsResult(self.collect_unused_references, [])
+
         try:
             namespace = await self.parent.documents_cache.get_namespace(document)
             if namespace is None:

robotcode/language_server/robotframework/parts/robocop_diagnostics.py

@@ -47,7 +47,7 @@ async def get_config(self, document: TextDocument) -> Optional[RoboCopConfig]:
     @language_id("robotframework")
     @threaded()
     @_logger.call
-    async def collect_diagnostics(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_diagnostics(self, sender: Any, document: TextDocument, full: bool) -> DiagnosticsResult:
 
         workspace_folder = self.parent.workspace.get_workspace_folder(document.uri)
         if workspace_folder is not None:

robotcode/utils/async_tools.py

@@ -494,7 +494,7 @@ def run() -> _T:
             try:
                 running_tasks = asyncio.all_tasks(loop)
                 if running_tasks:
-                    loop.run_until_complete(asyncio.gather(*running_tasks, loop=loop, return_exceptions=True))
+                    loop.run_until_complete(asyncio.gather(*running_tasks, return_exceptions=True))
 
                 loop.run_until_complete(loop.shutdown_asyncgens())
             finally:
@@ -601,27 +601,25 @@ def __repr__(self) -> str:
         return f"<{res[1:-1]} [{extra}]>"
 
     def _wake_up_next(self) -> None:
-        while self._waiters:
-            with self._lock:
+        with self._lock:
+            while self._waiters:
                 waiter = self._waiters.popleft()
 
-            if not waiter.done():
-                if waiter._loop == asyncio.get_running_loop():
-                    if not waiter.done():
-                        waiter.set_result(True)
-                else:
-                    if waiter._loop.is_running():
-
-                        def s(w: asyncio.Future[Any]) -> None:
-                            if not w.done():
-                                w.set_result(True)
-
-                        if not waiter.done():
-                            waiter._loop.call_soon_threadsafe(s, waiter)
-                    else:
-                        warnings.warn("Loop is not running.")
+                if not waiter.done():
+                    if waiter._loop == asyncio.get_running_loop():
                         if not waiter.done():
                             waiter.set_result(True)
+                    else:
+                        if waiter._loop.is_running():
+
+                            def s(w: asyncio.Future[Any]) -> None:
+                                if w._loop.is_running() and not w.done():
+                                    w.set_result(True)
+
+                            if not waiter.done():
+                                waiter._loop.call_soon_threadsafe(s, waiter)
+                        else:
+                            warnings.warn("Loop is not running.")
 
     def locked(self) -> bool:
         with self._lock:
@@ -688,6 +686,10 @@ async def acquire(self, timeout: Optional[float] = None) -> bool:
     def release(self) -> None:
         self._block.release()
 
+    @property
+    def locked(self) -> bool:
+        return self._block.locked()
+
     async def __aenter__(self) -> None:
         await self.acquire()
 
@@ -734,6 +736,10 @@ async def __inner_lock(self) -> AsyncGenerator[Any, None]:
         # with self._lock:
         #     yield None
 
+    @property
+    def locked(self) -> bool:
+        return self._locked
+
     async def acquire(self) -> bool:
         async with self.__inner_lock():
             if not self._locked and (self._waiters is None or all(w.cancelled() for w in self._waiters)):