refactor: introduce asyncio.RLock

Author: d-biehl

robotcode/language_server/common/parts/diagnostics.py

@@ -376,6 +376,7 @@ async def run_workspace_diagnostics(self) -> None:
                     f"collecting workspace diagnostics for for {len(documents)} "
                     f"documents takes {time.monotonic() - start}s"
                 )
+                self._logger.info(f"{len(self.parent.documents)} documents loaded")
 
             except (SystemExit, KeyboardInterrupt, asyncio.CancelledError):
                 raise
@@ -416,8 +417,6 @@ async def cancel(t: asyncio.Task[Any]) -> None:
                 except TimeoutError as e:
                     raise RuntimeError("Can't cancel diagnostics task.") from e
 
-                # task.get_loop().call_soon_threadsafe(task.cancel)
-
             data.version = document.version
             data.task = create_sub_task(
                 self._get_diagnostics_for_document(document, data, debounce, send_diagnostics),
@@ -459,18 +458,17 @@ async def _get_diagnostics_for_document(
                     if result.diagnostics is not None:
                         collected_keys.append(result.key)
 
-                    if data.entries:
-                        if send_diagnostics:
-                            self.parent.send_notification(
-                                "textDocument/publishDiagnostics",
-                                PublishDiagnosticsParams(
-                                    uri=document.document_uri,
-                                    version=document._version,
-                                    diagnostics=[
-                                        e for e in itertools.chain(*(i for i in data.entries.values() if i is not None))
-                                    ],
-                                ),
-                            )
+                    if data.entries and send_diagnostics:
+                        self.parent.send_notification(
+                            "textDocument/publishDiagnostics",
+                            PublishDiagnosticsParams(
+                                uri=document.document_uri,
+                                version=document._version,
+                                diagnostics=[
+                                    e for e in itertools.chain(*(i for i in data.entries.values() if i is not None))
+                                ],
+                            ),
+                        )
 
         except asyncio.CancelledError:
             self._logger.debug(lambda: f"_get_diagnostics cancelled for {document}")

robotcode/language_server/robotframework/parts/discovering.py

@@ -238,7 +238,6 @@ async def generate(suite: TestSuite) -> TestItem:
                 else None,
             )
 
-        # await self.parent.diagnostics.ensure_workspace_loaded()
         await self.parent.robot_workspace.documents_loaded.wait()
 
         start = time.monotonic()
@@ -247,107 +246,107 @@ async def generate(suite: TestSuite) -> TestItem:
 
         workspace_path = Uri(workspace_folder).to_path()
         canceled = False
-        with LOGGER.cache_only:
-            try:
-                config = await self.get_config(workspace_folder)
-                rpa_mode = config.get_rpa_mode() if config is not None else None
-                languages = await self.parent.documents_cache.get_workspace_languages(workspace_folder)
+        LOGGER._cache_only = True
+        try:
+            config = await self.get_config(workspace_folder)
+            rpa_mode = config.get_rpa_mode() if config is not None else None
+            languages = await self.parent.documents_cache.get_workspace_languages(workspace_folder)
 
-                if paths is None and config is not None:
-                    paths = config.paths
+            if paths is None and config is not None:
+                paths = config.paths
 
-                if paths and len(paths):
+            if paths and len(paths):
 
-                    def normalize_paths(paths: List[str]) -> Iterator[str]:
+                def normalize_paths(paths: List[str]) -> Iterator[str]:
 
-                        for path in paths:
+                    for path in paths:
 
-                            p = Path(path)
+                        p = Path(path)
 
-                            if not p.is_absolute():
-                                p = Path(workspace_path, p)
+                        if not p.is_absolute():
+                            p = Path(workspace_path, p)
 
-                            if p.exists():
-                                yield str(p)
+                        if p.exists():
+                            yield str(p)
 
-                    def nonexisting_paths(paths: List[str]) -> Iterator[str]:
+                def nonexisting_paths(paths: List[str]) -> Iterator[str]:
 
-                        for path in paths:
+                    for path in paths:
 
-                            p = Path(path)
+                        p = Path(path)
 
-                            if not p.is_absolute():
-                                p = Path(workspace_path, p)
+                        if not p.is_absolute():
+                            p = Path(workspace_path, p)
 
-                            if not p.exists():
-                                yield str(p)
+                        if not p.exists():
+                            yield str(p)
 
-                    valid_paths = [i for i in normalize_paths(paths)]
+                valid_paths = [i for i in normalize_paths(paths)]
 
-                    if get_robot_version() >= (6, 0):
-                        builder = TestSuiteBuilder(
-                            included_suites=suites if suites else None,
-                            rpa=rpa_mode,
-                            lang=languages,
-                        )
-                    else:
-                        builder = TestSuiteBuilder(included_suites=suites if suites else None, rpa=rpa_mode)
+                if get_robot_version() >= (6, 0):
+                    builder = TestSuiteBuilder(
+                        included_suites=suites if suites else None,
+                        rpa=rpa_mode,
+                        lang=languages,
+                    )
+                else:
+                    builder = TestSuiteBuilder(included_suites=suites if suites else None, rpa=rpa_mode)
 
-                    suite: Optional[TestSuite] = builder.build(*valid_paths) if valid_paths else None
-                    suite_item = [await generate(suite)] if suite else []
+                suite: Optional[TestSuite] = builder.build(*valid_paths) if valid_paths else None
+                suite_item = [await generate(suite)] if suite else []
 
-                    return [
-                        TestItem(
-                            type="workspace",
-                            id=str(Path.cwd().resolve()),
-                            label=Path.cwd().name,
-                            longname=Path.cwd().name,
-                            uri=str(Uri.from_path(Path.cwd())),
-                            children=[
-                                *suite_item,
-                                *[
-                                    TestItem(
-                                        type="error",
-                                        id=f"{i};ERROR",
-                                        longname="error",
-                                        label=i,
-                                        error=f"Parsing '{i}' failed: File or directory to does not exist.",
-                                    )
-                                    for i in nonexisting_paths(paths)
-                                ],
-                            ],
-                        )
-                    ]
-                else:
-                    if get_robot_version() >= (6, 0):
-                        builder = TestSuiteBuilder(
-                            included_suites=suites if suites else None,
-                            rpa=rpa_mode,
-                            lang=languages,
-                        )
-                    else:
-                        builder = TestSuiteBuilder(included_suites=suites if suites else None, rpa=rpa_mode)
-                    return [await generate(builder.build(str(workspace_path)))]
-            except (SystemExit, KeyboardInterrupt):
-                raise
-            except asyncio.CancelledError:
-                canceled = True
-                self._logger.info("Tests discovery canceled")
-                raise
-            except BaseException as e:
-                self._logger.info(f"Failed to discover tests: {e}")
                 return [
                     TestItem(
-                        type="error",
-                        id=str(Uri.from_path(Path.cwd().resolve())),
-                        longname="error",
+                        type="workspace",
+                        id=str(Path.cwd().resolve()),
                         label=Path.cwd().name,
-                        error=str(e),
+                        longname=Path.cwd().name,
+                        uri=str(Uri.from_path(Path.cwd())),
+                        children=[
+                            *suite_item,
+                            *[
+                                TestItem(
+                                    type="error",
+                                    id=f"{i};ERROR",
+                                    longname="error",
+                                    label=i,
+                                    error=f"Parsing '{i}' failed: File or directory to does not exist.",
+                                )
+                                for i in nonexisting_paths(paths)
+                            ],
+                        ],
                     )
                 ]
-            finally:
-                if not canceled:
-                    self._logger.info(f"Tests discovery took {time.monotonic() - start}s")
+            else:
+                if get_robot_version() >= (6, 0):
+                    builder = TestSuiteBuilder(
+                        included_suites=suites if suites else None,
+                        rpa=rpa_mode,
+                        lang=languages,
+                    )
+                else:
+                    builder = TestSuiteBuilder(included_suites=suites if suites else None, rpa=rpa_mode)
+                return [await generate(builder.build(str(workspace_path)))]
+        except (SystemExit, KeyboardInterrupt):
+            raise
+        except asyncio.CancelledError:
+            canceled = True
+            self._logger.info("Tests discovery canceled")
+            raise
+        except BaseException as e:
+            self._logger.info(f"Failed to discover tests: {e}")
+            return [
+                TestItem(
+                    type="error",
+                    id=str(Uri.from_path(Path.cwd().resolve())),
+                    longname="error",
+                    label=Path.cwd().name,
+                    error=str(e),
+                )
+            ]
+        finally:
+            if not canceled:
+                self._logger.info(f"Tests discovery took {time.monotonic() - start}s")
 
     @rpc_method(name="robot/discovering/getTestsFromDocument", param_type=GetTestsParams)
     @threaded()

robotcode/utils/async_cache.py

@@ -1,7 +1,7 @@
 from collections import defaultdict
 from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, cast
 
-from .async_tools import Lock
+from .async_tools import RLock
 
 _T = TypeVar("_T")
 
@@ -16,7 +16,7 @@ class CacheEntry:
     def __init__(self) -> None:
         self.data: Any = None
         self.has_data: bool = False
-        self.lock = Lock()
+        self.lock = RLock()
 
 
 class AsyncSimpleLRUCache:
@@ -27,7 +27,7 @@ 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()
+        self._lock = RLock()
 
     async def has(self, *args: Any, **kwargs: Any) -> bool:
         key = self._make_key(*args, **kwargs)

robotcode/utils/async_tools.py

@@ -466,12 +466,11 @@ def aaa(fut: asyncio.Future[Any]) -> None:
                     warnings.warn(f"Lock {self} takes to long {threading.current_thread()}\n, try to cancel...")
                     fut.cancel()
 
-                # h = fut.get_loop().call_later(120, aaa, "".join(traceback.format_stack()))
                 h = fut.get_loop().call_later(120, aaa, fut)
-
-                await fut
-
-                h.cancel()
+                try:
+                    await fut
+                finally:
+                    h.cancel()
             finally:
                 with self._lock:
                     self._waiters.remove(fut)
@@ -541,6 +540,29 @@ def set_result(w: asyncio.Future[Any], ev: threading.Event) -> None:
             self._wake_up_first()
 
 
+class RLock(Lock):
+    def __init__(self) -> None:
+        super().__init__()
+        self._task: Optional[asyncio.Task[Any]] = None
+        self._depth = 0
+
+    async def acquire(self) -> bool:
+        if self._task is None or self._task != asyncio.current_task():
+            await super().acquire()
+            self._task = asyncio.current_task()
+            assert self._depth == 0
+        self._depth += 1
+
+        return True
+
+    def release(self) -> None:
+        if self._depth > 0:
+            self._depth -= 1
+        if self._depth == 0:
+            super().release()
+            self._task = None
+
+
 class FutureInfo:
     def __init__(self, future: asyncio.Future[Any]) -> None:
         self.task: weakref.ref[asyncio.Future[Any]] = weakref.ref(future)