more speedups in collecting/finding references

Author: d-biehl

robotcode/language_server/common/parts/documents.py

@@ -69,6 +69,10 @@ def __init__(self, parent: LanguageServerProtocol) -> None:
         self.parent.on_initialized.add(self._protocol_initialized)
         self._lock = Lock()
 
+    @property
+    def documents(self) -> List[TextDocument]:
+        return list(self._documents.values())
+
     async def _protocol_initialized(self, sender: Any) -> None:
         await self._update_filewatchers()
 

robotcode/language_server/robotframework/parts/codelens.py

@@ -105,7 +105,7 @@ async def resolve(self, sender: Any, code_lens: CodeLens) -> Optional[CodeLens]:
         name = code_lens.data["name"]
         line = code_lens.data["line"]
 
-        if self.parent.robot_workspace.workspace_loaded:
+        if self.parent.robot_workspace.workspace_loaded.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:
@@ -128,6 +128,7 @@ async def find_refs() -> None:
                             kw_doc,
                             include_declaration=False,
                         )
+
                         # await self.parent.robot_references.find_keyword_references(
                         #     document, kw_doc, include_declaration=False
                         # )

robotcode/language_server/robotframework/parts/references.py

@@ -15,21 +15,14 @@
 )
 
 from ....utils.async_cache import AsyncSimpleLRUCache
-from ....utils.async_tools import run_coroutine_in_thread, threaded
-from ....utils.glob_path import iter_files
+from ....utils.async_tools import threaded
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
 from ...common.decorators import language_id
 from ...common.lsp_types import Location, Position, ReferenceContext
 from ...common.text_document import TextDocument
-from ..configuration import WorkspaceConfig
 from ..diagnostics.entities import LocalVariableDefinition, VariableDefinition
-from ..diagnostics.library_doc import (
-    RESOURCE_FILE_EXTENSION,
-    ROBOT_FILE_EXTENSION,
-    KeywordDoc,
-    LibraryDoc,
-)
+from ..diagnostics.library_doc import KeywordDoc, LibraryDoc
 from ..utils.ast_utils import (
     HasTokens,
     get_nodes_at_position,
@@ -108,36 +101,13 @@ async def _find_references_in_workspace(
         *args: Any,
         **kwargs: Any,
     ) -> List[Location]:
-        futures: List[Awaitable[List[Location]]] = []
+        await self.parent.robot_workspace.workspace_loaded.wait()
 
         result: List[Location] = []
 
-        for folder in self.parent.workspace.workspace_folders:
-
-            config = await self.parent.workspace.get_configuration(WorkspaceConfig, folder.uri) or WorkspaceConfig()
-
-            async for f in iter_files(
-                folder.uri.to_path(),
-                (f"**/*.{{{ROBOT_FILE_EXTENSION[1:]},{RESOURCE_FILE_EXTENSION[1:]}}}"),
-                ignore_patterns=config.exclude_patterns or [],  # type: ignore
-                absolute=True,
-            ):
-                try:
-                    doc = await self.parent.documents.get_or_open_document(f, "robotframework")
-                except (SystemExit, KeyboardInterrupt, asyncio.CancelledError):
-                    raise
-                except BaseException as ex:
-                    self._logger.exception(ex)
-                else:
-                    futures.append(run_coroutine_in_thread(func, doc, *args, **kwargs))
-
-            for e in await asyncio.gather(*futures, return_exceptions=True):
-                if isinstance(e, BaseException):
-                    if not isinstance(result, asyncio.CancelledError):
-                        self._logger.exception(e)
-                    continue
-                if e:
-                    result.extend(e)
+        for doc in self.parent.documents.documents:
+            if doc.language_id == "robotframework":
+                result.extend(await func(doc, *args, **kwargs))
 
         return result
 

robotcode/language_server/robotframework/parts/robot_workspace.py

@@ -3,7 +3,7 @@
 import asyncio
 from typing import TYPE_CHECKING, Any, List, Optional
 
-from ....utils.async_tools import threaded
+from ....utils.async_tools import Event, threaded
 from ....utils.glob_path import iter_files
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
@@ -30,7 +30,7 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         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.on_get_diagnostics_mode.add(self.on_get_diagnostics_mode)
-        self.workspace_loaded = False
+        self.workspace_loaded = Event()
 
     @language_id("robotframework")
     async def _on_read_document_text(self, sender: Any, uri: Uri) -> Optional[str]:
@@ -101,7 +101,7 @@ async def _collect_workspace_documents(self, sender: Any) -> List[WorkspaceDocum
                     except BaseException as e:
                         self._logger.exception(e)
 
-        self.workspace_loaded = True
+        self.workspace_loaded.set()
 
         if config.analysis.references_code_lens:
             await self.parent.code_lens.refresh()

robotcode/utils/async_tools.py

@@ -475,6 +475,7 @@ async def create_inner_task() -> _T:
                 inner_task.cancel()
 
             return await inner_task
+
         finally:
             threading.current_thread().setName(old_name)
 
@@ -585,6 +586,8 @@ async def __inner_lock(self) -> AsyncGenerator[Any, None]:
             yield None
         finally:
             self._lock.release()
+        # with self._lock:
+        #     yield None
 
     async def acquire(self) -> bool:
         async with self.__inner_lock():
@@ -603,7 +606,8 @@ async def acquire(self) -> bool:
                 await fut
             finally:
                 async with self.__inner_lock():
-                    self._waiters.remove(fut)
+                    if fut in self._waiters:
+                        self._waiters.remove(fut)
                     self._locked = True
         except asyncio.CancelledError:
             await self._wake_up_next()
@@ -621,16 +625,19 @@ async def release(self) -> None:
                 else:
                     warnings.warn(f"Lock is not acquired ({len(self._waiters) if self._waiters else 0} waiters).")
 
-            await self._wake_up_next()
+        await self._wake_up_next()
 
     async def _wake_up_next(self) -> None:
         if not self._waiters:
             return
 
-        try:
-            fut = next(iter(self._waiters))
-        except StopIteration:
-            return
+        async with self.__inner_lock():
+            try:
+                fut = next(iter(self._waiters))
+            except StopIteration:
+                return
+            if fut in self._waiters:
+                self._waiters.remove(fut)
 
         def s() -> None:
             if not fut.done():
@@ -640,7 +647,8 @@ def s() -> None:
             if not fut.done():
                 fut.set_result(True)
         else:
-            fut._loop.call_soon_threadsafe(s)
+            if fut._loop.is_running():
+                fut._loop.call_soon_threadsafe(s)
 
 
 class FutureInfo: