optimize namespace diagnostics and canceling diagnostics

Author: d-biehl

robotcode/language_server/common/parts/diagnostics.py

@@ -27,35 +27,41 @@ class PublishDiagnosticsEntry:
 
     def __init__(
         self,
-        document_uri: DocumentUri,
+        uri: Uri,
         cancelation_token: CancelationToken,
         task_factory: Callable[..., asyncio.Task[Any]],
+        done_callback: Callable[[PublishDiagnosticsEntry], Any],
     ) -> None:
 
-        self._document_uri = document_uri
+        self.uri = uri
 
         self._task_factory = task_factory
+        self.done_callback = done_callback
 
         self._task: Optional[asyncio.Task[Any]] = None
 
         self.cancel_token = cancelation_token
+        self.done = False
 
         @PublishDiagnosticsEntry._logger.call
         def create_task() -> None:
             self._task = self._task_factory()
 
             if self._task is not None:
-                self._task.set_name(f"Diagnostics for {self._document_uri}")
+                self._task.set_name(f"Diagnostics for {self.uri}")
 
                 def _done(t: asyncio.Task[Any]) -> None:
                     self._task = None
+                    self.done = True
+                    self.done_callback(self)
 
                 self._task.add_done_callback(_done)
 
         self._timer_handle: asyncio.TimerHandle = asyncio.get_event_loop().call_later(DIAGNOSTICS_DEBOUNCE, create_task)
 
     def __del__(self) -> None:
-        self.cancel(_from_del=True)
+        if self.task is not None:
+            self.cancel(_from_del=True)
 
     @property
     def task(self) -> Optional[asyncio.Task[Any]]:
@@ -65,14 +71,16 @@ def __str__(self) -> str:
         return self.__repr__()
 
     def __repr__(self) -> str:
-        return f"{type(self)}(document={repr(self._document_uri)}, task={repr(self.task)})"
+        return f"{type(self)}(document={repr(self.uri)}, task={repr(self.task)}, done={self.done})"
 
     @_logger.call(condition=lambda self, _from_del=False: not _from_del)
-    def cancel(self, *, _from_del: Optional[bool] = False) -> None:
+    def cancel(self, *, _from_del: Optional[bool] = False) -> Optional[asyncio.Task[None]]:
         self._timer_handle.cancel()
 
         if self.task is None:
-            return
+            return None
+
+        self.cancel_token.cancel()
 
         async def cancel() -> None:
             if self.task is None:
@@ -81,7 +89,6 @@ async def cancel() -> None:
             t = self.task
             self._task = None
             if not t.done():
-                self.cancel_token.cancel()
                 t.cancel()
                 try:
                     await t
@@ -90,7 +97,7 @@ async def cancel() -> None:
                 except BaseException:
                     pass
 
-        asyncio.create_task(cancel())
+        return asyncio.create_task(cancel())
 
 
 @dataclass
@@ -107,7 +114,6 @@ def __init__(self, protocol: LanguageServerProtocol) -> None:
 
         self._running_diagnostics: Dict[Uri, PublishDiagnosticsEntry] = {}
         self._task_lock = asyncio.Lock()
-        self._start_lock_lock = asyncio.Lock()
 
         self.parent.on_connection_lost.add(self.on_connection_lost)
         self.parent.on_shutdown.add(self.on_shutdown)
@@ -145,11 +151,13 @@ async def _cancel_all_tasks(self) -> None:
                 self._cancel_entry(v)
 
     @_logger.call(condition=lambda self, entry: entry is not None)
-    def _cancel_entry(self, entry: Optional[PublishDiagnosticsEntry]) -> None:
+    async def _cancel_entry(self, entry: Optional[PublishDiagnosticsEntry]) -> None:
         if entry is None:
             return
-
-        entry.cancel()
+        if not entry.done:
+            cancel_task = entry.cancel()
+            if cancel_task is not None:
+                await cancel_task
 
     @language_id("robotframework")
     @_logger.call
@@ -165,24 +173,30 @@ async def on_did_save(self, sender: Any, document: TextDocument) -> None:
     @_logger.call
     async def on_did_close(self, sender: Any, document: TextDocument) -> None:
         async with self._task_lock:
-            e = self._running_diagnostics.pop(document.uri, None)
-            self._cancel_entry(e)
+            self._cancel_entry(self._running_diagnostics.get(document.uri, None))
 
     @_logger.call
     async def on_did_change(self, sender: Any, document: TextDocument) -> None:
         await self.start_publish_diagnostics_task(document)
 
+    @_logger.call
+    def _delete_entry(self, e: PublishDiagnosticsEntry) -> None:
+        if self._running_diagnostics[e.uri] == e:
+            self._running_diagnostics.pop(e.uri, None)
+
     @_logger.call
     async def start_publish_diagnostics_task(self, document: TextDocument) -> None:
+        await self._cancel_entry(self._running_diagnostics.get(document.uri, None))
+
         async with self._task_lock:
-            self._cancel_entry(self._running_diagnostics.get(document.uri, None))
             cancelation_token = CancelationToken()
             self._running_diagnostics[document.uri] = PublishDiagnosticsEntry(
-                document.document_uri,
+                document.uri,
                 cancelation_token,
                 lambda: asyncio.create_task(
                     self.publish_diagnostics(document.document_uri, cancelation_token),
                 ),
+                self._delete_entry,
             )
 
     @_logger.call
@@ -202,11 +216,14 @@ async def publish_diagnostics(self, document_uri: DocumentUri, cancelation_token
             callback_filter=language_id_filter(document),
             return_exceptions=True,
         ):
+            if cancelation_token.canceled:
+                break
+
             result = cast(DiagnosticsResult, result_any)
 
             if isinstance(result, BaseException):
-                if not isinstance(result, asyncio.CancelledError):
-                    self._logger.exception(result, exc_info=result)
+                # if not isinstance(result, asyncio.CancelledError):
+                self._logger.exception(result, exc_info=result)
             else:
 
                 diagnostics[result.key] = result.diagnostics if result.diagnostics else []

robotcode/language_server/robotframework/diagnostics/namespace.py

@@ -573,6 +573,7 @@ async def invalidate(self) -> None:
             self._library_doc = None
             self._analyzed = False
             self._diagnostics = []
+            self._finder = None
 
         self.invalidated_callback(self)
 
@@ -631,6 +632,9 @@ async def ensure_initialized(self) -> bool:
         if not self._initialized:
             async with self._initialize_lock:
                 if not self._initialized:
+
+                    self._logger.debug(f"ensure_initialized -> initialize {self.document}")
+
                     imports = await self.get_imports()
 
                     data_entry: Optional[Namespace.DataEntry] = None
@@ -675,8 +679,11 @@ async def ensure_initialized(self) -> bool:
 
                     self._initialized = True
 
-            if not self._initialized:
-                raise Exception("Namespace not initialized")
+                    await self.get_library_doc()
+                    await self.get_libraries_matchers()
+                    await self.get_resources_matchers()
+                    await self.get_finder()
+
         return self._initialized
 
     @property
@@ -1096,9 +1103,12 @@ async def _analyze(self, cancelation_token: Optional[CancelationToken] = None) -
                 if not self._analyzed:
                     canceled = False
                     try:
-                        self._diagnostics += await awaitable_run_in_thread(
-                            Analyzer().get(self.model, self, cancelation_token)
-                        )
+                        result = await awaitable_run_in_thread(Analyzer().get(self.model, self, cancelation_token))
+
+                        if cancelation_token is not None:
+                            cancelation_token.throw_if_canceled()
+
+                        self._diagnostics += result
 
                         lib_doc = await self.get_library_doc()
 
@@ -1124,24 +1134,23 @@ async def _analyze(self, cancelation_token: Optional[CancelationToken] = None) -
                                 )
                     except asyncio.CancelledError:
                         canceled = True
+                        self._logger.debug("analyzing canceled")
                         raise
                     finally:
                         self._analyzed = not canceled
+                        self._logger.debug(
+                            f"analyzed {self.document}" if self._analyzed else f"not analyzed {self.document}"
+                        )
 
-    @_logger.call(condition=lambda self, name: self._finder is not None and name not in self._finder._cache)
-    async def find_keyword(self, name: Optional[str]) -> Optional[KeywordDoc]:
+    async def get_finder(self) -> KeywordFinder:
         if self._finder is None:
             await self.ensure_initialized()
+            self._finder = KeywordFinder(self)
+        return self._finder
 
-            self._finder = await self.create_finder()
-
-        return await self._finder.find_keyword(name)
-
-    @_logger.call
-    async def create_finder(self) -> KeywordFinder:
-        await self.ensure_initialized()
-
-        return KeywordFinder(self)
+    @_logger.call(condition=lambda self, name: self._finder is not None and name not in self._finder._cache)
+    async def find_keyword(self, name: Optional[str]) -> Optional[KeywordDoc]:
+        return await (await self.get_finder()).find_keyword(name)
 
 
 class DiagnosticsEntry(NamedTuple):

robotcode/language_server/robotframework/parts/diagnostics.py

@@ -159,4 +159,6 @@ async def collect_namespace_diagnostics(
         if namespace is None:
             return DiagnosticsResult(self.collect_namespace_diagnostics, None)
 
+        await namespace.ensure_initialized()
+
         return DiagnosticsResult(self.collect_namespace_diagnostics, await namespace.get_diagnostisc(cancelation_token))

robotcode/language_server/robotframework/parts/robocop_diagnostics.py

@@ -92,6 +92,7 @@ def collect(
         from robocop.config import Config
         from robocop.rules import RuleSeverity
         from robocop.run import Robocop
+        from robocop.utils.misc import is_suite_templated
 
         result: List[Diagnostic] = []
 
@@ -109,11 +110,30 @@ def collect(
             if extension_config.configurations:
                 config.configure = set(extension_config.configurations)
 
-            analyser = Robocop(from_cli=False, config=config)
+            class MyRobocop(Robocop):  # type: ignore
+                def run_check(self, ast_model, filename, source=None):  # type: ignore
+                    found_issues = []
+                    self.register_disablers(filename, source)
+                    if self.disabler.file_disabled:
+                        return []
+                    templated = is_suite_templated(ast_model)
+                    for checker in self.checkers:
+                        cancelation_token.throw_if_canceled()
+
+                        if checker.disabled:
+                            continue
+                        found_issues += [
+                            issue
+                            for issue in checker.scan_file(ast_model, filename, source, templated)
+                            if not self.disabler.is_rule_disabled(issue)
+                        ]
+                    return found_issues
+
+            analyser = MyRobocop(from_cli=False, config=config)
             analyser.reload_config()
 
             # TODO find a way to cancel the run_check
-            issues = analyser.run_check(model, str(document.uri.to_path()), document.text)
+            issues = analyser.run_check(model, str(document.uri.to_path()), document.text)  # type: ignore
 
             for issue in issues:
                 d = Diagnostic(

robotcode/language_server/robotframework/parts/semantic_tokens.py

@@ -627,6 +627,7 @@ async def collect_threading(
         try:
             model = await self.parent.documents_cache.get_model(document)
             namespace = await self.parent.documents_cache.get_namespace(document)
+            await namespace.ensure_initialized()
 
             builtin_library_doc = next(
                 (

robotcode/utils/async_tools.py

@@ -397,17 +397,14 @@ def __init__(self, _func: _TCallable) -> None:
 
 class CancelationToken:
     def __init__(self) -> None:
-        self._canceled = False
-        self._lock = threading.RLock()
+        self._canceled = asyncio.Event()
 
     @property
     def canceled(self) -> bool:
-        with self._lock:
-            return self._canceled
+        return self._canceled.is_set()
 
     def cancel(self) -> None:
-        with self._lock:
-            self._canceled = True
+        self._canceled.set()
 
     def throw_if_canceled(self) -> bool:
         if self.canceled: