stabilize document diagnostics

Author: d-biehl

robotcode/jsonrpc2/server.py

@@ -79,6 +79,8 @@ def __init__(
         self._closed = False
 
         self.loop = asyncio.get_event_loop()
+        if self.loop is not None:
+            self.loop.slow_callback_duration = 10
 
     def __del__(self) -> None:
         self.close()

robotcode/language_server/common/parts/diagnostics.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import itertools
+import threading
 import uuid
 from dataclasses import dataclass, field
 from enum import Enum
@@ -74,10 +75,34 @@ class WorkspaceDocumentsResult:
 
 @dataclass
 class DiagnosticsData:
-    id: str
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
     entries: Dict[Any, Optional[List[Diagnostic]]] = field(default_factory=dict)
     version: Optional[int] = None
     task: Optional[asyncio.Task[Any]] = None
+    force: bool = False
+
+
+def _cancel_all_tasks(loop: asyncio.AbstractEventLoop) -> None:
+    to_cancel = asyncio.all_tasks(loop)
+    if not to_cancel:
+        return
+
+    for task in to_cancel:
+        task.cancel()
+
+    loop.run_until_complete(asyncio.gather(*to_cancel, loop=loop, return_exceptions=True))
+
+    for task in to_cancel:
+        if task.cancelled():
+            continue
+        if task.exception() is not None:
+            loop.call_exception_handler(
+                {
+                    "message": "unhandled exception during asyncio.run() shutdown",
+                    "exception": task.exception(),
+                    "task": task,
+                }
+            )
 
 
 class DiagnosticsProtocolPart(LanguageServerProtocolPart, HasExtendCapabilities):
@@ -97,18 +122,63 @@ def __init__(self, protocol: LanguageServerProtocol) -> None:
 
         self._current_workspace_task: Optional[asyncio.Task[WorkspaceDiagnosticReport]] = None
 
+        self._diagnostics_loop: Optional[asyncio.AbstractEventLoop] = None
+        self._diagnostics_loop_lock = threading.RLock()
+        self._diagnostics_started = threading.Event()
+
+        self.parent.on_initialized.add(self.initialized)
+
         self.in_get_workspace_diagnostics = Event(True)
 
+    async def initialized(self, sender: Any) -> None:
+        self._ensure_diagnostics_thread_started()
+
+    @property
+    def diagnostics_loop(self) -> asyncio.AbstractEventLoop:
+        if self._diagnostics_loop is None:
+            self._ensure_diagnostics_thread_started()
+
+        assert self._diagnostics_loop is not None
+
+        return self._diagnostics_loop
+
+    def _run_diagnostics(self) -> None:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            self._diagnostics_loop = loop
+            self._diagnostics_started.set()
+
+            loop.slow_callback_duration = 10
+
+            loop.run_forever()
+            _cancel_all_tasks(loop)
+            loop.run_until_complete(loop.shutdown_asyncgens())
+        finally:
+            asyncio.set_event_loop(None)
+            loop.close()
+
+    def _ensure_diagnostics_thread_started(self) -> None:
+        with self._diagnostics_loop_lock:
+            if self._diagnostics_loop is None:
+                self._server_thread = threading.Thread(
+                    name="diagnostics_worker", target=self._run_diagnostics, daemon=True
+                )
+
+                self._server_thread.start()
+
+                if not self._diagnostics_started.wait(10):
+                    raise RuntimeError("Can't start diagnostics worker thread.")
+
     def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         if (
             self.parent.client_capabilities is not None
             and self.parent.client_capabilities.text_document is not None
             and self.parent.client_capabilities.text_document.diagnostic is not None
         ):
-            # capabilities.diagnostic_provider = None
             capabilities.diagnostic_provider = DiagnosticOptions(
                 inter_file_dependencies=True,
-                workspace_diagnostics=True,
+                workspace_diagnostics=False,
                 identifier=f"robotcodelsp_{uuid.uuid4()}",
                 work_done_progress=True,
             )
@@ -117,19 +187,6 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
     async def collect(sender, document: TextDocument) -> DiagnosticsResult:  # NOSONAR
         ...
 
-    @async_tasking_event_iterator
-    async def collect_document_has_diagnostics(sender, document: TextDocument) -> bool:  # NOSONAR
-        ...
-
-    async def document_has_diagnostics(self, document: TextDocument) -> bool:  # NOSONAR
-        async for result in self.collect_document_has_diagnostics(
-            self, document, callback_filter=language_id_filter(document)
-        ):
-            if result:
-                return True
-
-        return False
-
     @async_tasking_event
     async def load_workspace_documents(sender) -> List[WorkspaceDocumentsResult]:  # NOSONAR
         ...
@@ -167,10 +224,22 @@ async def ensure_workspace_loaded(self) -> None:
                         self._workspace_loaded = True
                         self.workspace_loaded_event.set()
                         await self.on_workspace_loaded(self)
-                        await self.refresh()
+                        await self.force_refresh_all()
+
+    async def force_refresh_all(self) -> None:
+        for doc in self.parent.documents.documents:
+            self.get_diagnostics_data(doc).force = True
+
+        await self.refresh()
+
+    async def force_refresh_document(self, document: TextDocument) -> None:
+        self.get_diagnostics_data(document).force = True
+        if document.opened_in_editor:
+            await self.refresh()
 
     @_logger.call
-    async def _get_diagnostics(self, document: TextDocument, data: DiagnosticsData) -> None:
+    async def _get_diagnostics_for_document(self, document: TextDocument, data: DiagnosticsData) -> None:
+        self._logger.debug(lambda: f"Get diagnostics for {document}")
 
         await asyncio.sleep(0.75)
 
@@ -193,12 +262,11 @@ async def _get_diagnostics(self, document: TextDocument, data: DiagnosticsData)
                     if result.diagnostics is not None:
                         collected_keys.append(result.key)
 
-                    await self.refresh()
+                    if document.opened_in_editor:
+                        await self.refresh()
 
         except asyncio.CancelledError:
-            self._logger.critical(lambda: f"_get_diagnostics cancelled for {document}")
-        else:
-            await self.refresh()
+            self._logger.debug(lambda: f"_get_diagnostics cancelled for {document}")
         finally:
             for k in set(data.entries.keys()) - set(collected_keys):
                 data.entries.pop(k)
@@ -214,50 +282,41 @@ async def _text_document_diagnostic(
         **kwargs: Any,
     ) -> DocumentDiagnosticReport:
         try:
-            # if not self.workspace_loaded_event.is_set():
-            #     raise JsonRPCErrorException(
-            #         ErrorCodes.SERVER_CANCELLED,
-            #         "Workspace not loaded.",
-            #         data=DiagnosticServerCancellationData(True),
-            #     )
-
             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.")
+                raise JsonRPCErrorException(ErrorCodes.SERVER_CANCELLED, f"Document {text_document!r} not found.")
 
-            data: DiagnosticsData = document.get_data(self, None)
+            data = self.get_diagnostics_data(document)
 
-            if data is None:
-                data = DiagnosticsData(str(uuid.uuid4()))
-                document.set_data(self, data)
-
-            if (
-                document.version != data.version
-                or data.task is None
-                or not await self.document_has_diagnostics(document)
-            ):
+            if data.force or document.version != data.version or data.task is None:
 
                 task = data.task
 
-                data = DiagnosticsData(str(uuid.uuid4()))
+                data = DiagnosticsData()
                 document.set_data(self, data)
 
                 if task is not None and not task.done():
-                    self._logger.critical(lambda: f"try to cancel diagnostics for {document}")
+                    self._logger.debug(lambda: f"try to cancel diagnostics for {document}")
                     task.get_loop().call_soon_threadsafe(task.cancel)
 
                 data.version = document.version
                 data.task = create_sub_task(
-                    self._get_diagnostics(document, data), loop=self.parent.loop, name=f"diagnostics ${text_document}"
+                    self._get_diagnostics_for_document(document, data),
+                    loop=self.diagnostics_loop,
+                    name=f"diagnostics ${text_document}",
                 )
 
                 def done(t: asyncio.Task[Any]) -> None:
-                    if t.cancelled():
-                        self._logger.critical(lambda: f"diagnostics for {document} canceled")
-                        try:
-                            t.exception()
-                        except asyncio.CancelledError:
-                            pass
+
+                    self._logger.debug(lambda: f"diagnostics for {document} {'canceled' if t.cancelled() else 'ended'}")
+                    try:
+                        t.result()
+                    except asyncio.CancelledError:
+                        pass
+                    except (SystemExit, KeyboardInterrupt):
+                        raise
+                    except BaseException as e:
+                        self._logger.exception(e)
 
                 data.task.add_done_callback(done)
 
@@ -268,9 +327,18 @@ def done(t: asyncio.Task[Any]) -> None:
                 list(itertools.chain(*(e for e in data.entries.values() if e is not None))), result_id=data.id
             )
         except asyncio.CancelledError:
-            self._logger.critical("canceled _text_document_diagnostic")
+            self._logger.debug("canceled _text_document_diagnostic")
             raise
 
+    def get_diagnostics_data(self, document: TextDocument) -> DiagnosticsData:
+        data: DiagnosticsData = document.get_data(self, None)
+
+        if data is None:
+            data = DiagnosticsData(str(uuid.uuid4()))
+            document.set_data(self, data)
+
+        return data
+
     @rpc_method(name="workspace/diagnostic", param_type=WorkspaceDiagnosticParams)
     @threaded()
     async def _workspace_diagnostic(

robotcode/language_server/common/parts/semantic_tokens.py

@@ -173,4 +173,10 @@ async def _text_document_semantic_tokens_range(
         return None
 
     async def workspace_refresh(self) -> None:
-        await self.parent.send_request_async("workspace/semanticTokens/refresh")
+        if (
+            self.parent.client_capabilities is not None
+            and self.parent.client_capabilities.workspace is not None
+            and self.parent.client_capabilities.workspace.semantic_tokens is not None
+            and self.parent.client_capabilities.workspace.semantic_tokens.refresh_support
+        ):
+            await self.parent.send_request_async("workspace/semanticTokens/refresh")

robotcode/language_server/robotframework/diagnostics/namespace.py

@@ -342,8 +342,6 @@ async def visit_InlineIfHeader(self, node: ast.AST) -> None:  # noqa: N802
         from robot.parsing.lexer.tokens import Token as RobotToken
         from robot.parsing.model.statements import InlineIfHeader
 
-        # TODO  analyze "Set Local/Global/Suite Variable"
-
         n = cast(InlineIfHeader, node)
 
         for assign_token in n.get_tokens(RobotToken.ASSIGN):

robotcode/language_server/robotframework/parts/code_action.py

@@ -166,7 +166,7 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
 
         self._documentation_server: Optional[ThreadingHTTPServer] = None
         self._documentation_server_lock = threading.RLock()
-        self._documentation_server_port = check_free_port(3000)
+        self._documentation_server_port = 0
 
     async def initialized(self, sender: Any) -> None:
         self._ensure_http_server_started()
@@ -178,6 +178,7 @@ async def shutdown(self, sender: Any) -> None:
                 self._documentation_server = None
 
     def _run_server(self) -> None:
+        self._documentation_server_port = check_free_port(3100)
         with DualStackServer(("", self._documentation_server_port), LibDocRequestHandler) as server:
             self._documentation_server = server
             try:

robotcode/language_server/robotframework/parts/diagnostics.py

@@ -48,13 +48,12 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         parent.diagnostics.collect.add(self.collect_unused_keyword_references)
         parent.diagnostics.collect.add(self.collect_unused_variable_references)
 
-        parent.diagnostics.collect_document_has_diagnostics.add(self.collect_document_has_diagnostics)
+        parent.documents_cache.namespace_invalidated.add(self.namespace_invalidated)
 
     @language_id("robotframework")
     @_logger.call
-    async def collect_document_has_diagnostics(self, sender: Any, document: TextDocument) -> bool:
-        namespace = await self.parent.documents_cache.get_namespace(document)
-        return namespace is None or await namespace.is_initialized() and await namespace.is_analyzed()
+    async def namespace_invalidated(self, sender: Any, document: TextDocument) -> None:
+        await self.parent.diagnostics.force_refresh_document(document)
 
     @language_id("robotframework")
     @threaded()
@@ -311,20 +310,20 @@ async def _collect_unused_keyword_references(self, document: TextDocument) -> Di
     @_logger.call
     async def collect_unused_variable_references(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
         if not self.parent.diagnostics.workspace_loaded_event.is_set():
-            return DiagnosticsResult(self.collect_unused_keyword_references, None)
+            return DiagnosticsResult(self.collect_unused_variable_references, None)
 
         config = await self.parent.workspace.get_configuration(AnalysisConfig, document.uri)
 
         if not config.find_unused_references:
-            return DiagnosticsResult(self.collect_unused_keyword_references, [])
+            return DiagnosticsResult(self.collect_unused_variable_references, [])
 
         return await self._collect_unused_variable_references(document)
 
     async def _collect_unused_variable_references(self, document: TextDocument) -> DiagnosticsResult:
         try:
             namespace = await self.parent.documents_cache.get_namespace(document)
             if namespace is None:
-                return DiagnosticsResult(self.collect_unused_keyword_references, None)
+                return DiagnosticsResult(self.collect_unused_variable_references, None)
 
             result: List[Diagnostic] = []
 
@@ -343,12 +342,12 @@ async def _collect_unused_variable_references(self, document: TextDocument) -> D
                         )
                     )
 
-            return DiagnosticsResult(self.collect_unused_keyword_references, result)
+            return DiagnosticsResult(self.collect_unused_variable_references, result)
         except (asyncio.CancelledError, SystemExit, KeyboardInterrupt):
             raise
         except BaseException as e:
             return DiagnosticsResult(
-                self.collect_unused_keyword_references,
+                self.collect_unused_variable_references,
                 [
                     Diagnostic(
                         range=Range(

robotcode/language_server/robotframework/parts/semantic_tokens.py

@@ -137,6 +137,13 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         # parent.semantic_tokens.collect_range.add(self.collect_range)
         # parent.semantic_tokens.collect_full_delta.add(self.collect_full_delta)
 
+        parent.documents_cache.namespace_invalidated.add(self.namespace_invalidated)
+
+    @language_id("robotframework")
+    @_logger.call
+    async def namespace_invalidated(self, sender: Any, document: TextDocument) -> None:
+        await self.parent.semantic_tokens.workspace_refresh()
+
     @classmethod
     def generate_mapping(cls) -> Dict[str, Tuple[Enum, Optional[Set[Enum]]]]:
         from robot.parsing.lexer.tokens import Token as RobotToken