implement threaded rpc_methods

Author: d-biehl

robotcode/jsonrpc2/protocol.py

@@ -12,7 +12,6 @@
 from typing import (
     Any,
     Callable,
-    Coroutine,
     Dict,
     Generic,
     Iterator,
@@ -477,23 +476,27 @@ def _handle_body(self, body: bytes, charset: str) -> None:
 
     def _handle_messages(self, iterator: Iterator[JsonRPCMessage]) -> None:
         def done(f: asyncio.Future[Any]) -> None:
-            ex = f.exception()
-            if ex is not None and not isinstance(ex, asyncio.CancelledError):
-                self._logger.exception(ex, exc_info=ex)
+            if f.done() and not f.cancelled():
+                ex = f.exception()
+
+                if ex is None or isinstance(ex, asyncio.CancelledError):
+                    return
+
+                # self._logger.exception(ex, exc_info=ex)
 
         for m in iterator:
             create_sub_task(self.handle_message(m)).add_done_callback(done)
 
     @_logger.call
     async def handle_message(self, message: JsonRPCMessage) -> None:
         if isinstance(message, JsonRPCRequest):
-            self.handle_request(message)
+            await self.handle_request(message)
         elif isinstance(message, JsonRPCNotification):
-            self.handle_notification(message)
+            await self.handle_notification(message)
         elif isinstance(message, JsonRPCError):
-            self.handle_error(message)
+            await self.handle_error(message)
         elif isinstance(message, JsonRPCResponse):
-            self.handle_response(message)
+            await self.handle_response(message)
 
     @_logger.call
     def send_response(self, id: Optional[Union[str, int, None]], result: Optional[Any] = None) -> None:
@@ -570,7 +573,7 @@ def send_notification(self, method: str, params: Any) -> None:
         self.send_message(JsonRPCNotification(method=method, params=params))
 
     @_logger.call(exception=True)
-    def handle_response(self, message: JsonRPCResponse) -> None:
+    async def handle_response(self, message: JsonRPCResponse) -> None:
         if message.id is None:
             error = "Invalid response. Response id is null."
             self._logger.warning(error)
@@ -614,7 +617,7 @@ def s(f: asyncio.Future[Any], r: Any) -> None:
                     entry.future._loop.call_soon_threadsafe(entry.future.set_exception, e)
 
     @_logger.call
-    def handle_error(self, message: JsonRPCError) -> None:
+    async def handle_error(self, message: JsonRPCError) -> None:
         raise JsonRPCErrorException(message.error.code, message.error.message, message.error.data)
 
     @staticmethod
@@ -674,7 +677,7 @@ def _convert_params(
         return args, kw_args
 
     @_logger.call
-    def handle_request(self, message: JsonRPCRequest) -> Optional[asyncio.Task[_T]]:
+    async def handle_request(self, message: JsonRPCRequest) -> None:
         e = self.registry.get_entry(message.method)
 
         if e is None or not callable(e.method):
@@ -716,7 +719,7 @@ def done(t: asyncio.Task[Any]) -> None:
 
         task.add_done_callback(done)
 
-        return task
+        await task
 
     @_logger.call
     def cancel_request(self, id: Union[int, str, None]) -> None:
@@ -734,18 +737,24 @@ async def cancel_all_received_request(self) -> None:
                 entry.future.cancel()
 
     @_logger.call
-    def handle_notification(self, message: JsonRPCNotification) -> None:
+    async def handle_notification(self, message: JsonRPCNotification) -> None:
         e = self.registry.get_entry(message.method)
 
         if e is None or not callable(e.method):
             self._logger.warning(f"Unknown method: {message.method}")
             return
         try:
             params = self._convert_params(e.method, e.param_type, message.params)
-            result = e.method(*params[0], **params[1])
-            if inspect.isawaitable(result):
-                create_sub_task(cast(Coroutine[Any, Any, Any], result))
 
+            if isinstance(e.method, HasThreaded) and cast(HasThreaded, e.method).__threaded__:
+                task = run_coroutine_in_thread(ensure_coroutine(e.method), *params[0], **params[1])
+            else:
+                task = create_sub_task(
+                    ensure_coroutine(e.method)(*params[0], **params[1]),
+                    name=message.method,
+                )
+
+            await task
         except asyncio.CancelledError:
             pass
         except (SystemExit, KeyboardInterrupt):

robotcode/language_server/common/parts/diagnostics.py

@@ -14,6 +14,7 @@
     async_tasking_event,
     async_tasking_event_iterator,
     create_sub_task,
+    threaded,
 )
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
@@ -101,7 +102,7 @@ def extend_capabilities(self, capabilities: ServerCapabilities) -> None:
         ):
             capabilities.diagnostic_provider = DiagnosticOptions(
                 inter_file_dependencies=True,
-                workspace_diagnostics=True,
+                workspace_diagnostics=False,
                 identifier=f"robotcodelsp_{uuid.uuid4()}",
                 work_done_progress=True,
             )
@@ -190,11 +191,7 @@ async def __get_full_document_diagnostics(self, document: TextDocument) -> Relat
         diagnostics: List[Diagnostic] = []
 
         async for result_any in self.collect(
-            self,
-            document,
-            full=True,
-            callback_filter=language_id_filter(document),
-            return_exceptions=True,
+            self, document, full=True, callback_filter=language_id_filter(document), return_exceptions=True
         ):
             result = cast(DiagnosticsResult, result_any)
 
@@ -207,6 +204,7 @@ async def __get_full_document_diagnostics(self, document: TextDocument) -> Relat
         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,
@@ -265,6 +263,7 @@ 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],

robotcode/utils/async_tools.py

@@ -2,11 +2,11 @@
 
 import asyncio
 import concurrent.futures
+import contextlib
 import contextvars
 import functools
 import inspect
 import threading
-import warnings
 import weakref
 from collections import deque
 from concurrent.futures.thread import ThreadPoolExecutor
@@ -523,6 +523,22 @@ def done(task: asyncio.Future[_T]) -> None:
     return result
 
 
+_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+
+
+@contextlib.asynccontextmanager
+async def async_lock(lock: threading.RLock) -> AsyncGenerator[None, None]:
+    locked = lock.acquire(False)
+    while not locked:
+        await asyncio.sleep(0)
+        locked = lock.acquire(False)
+    try:
+        yield
+    finally:
+        if locked:
+            lock.release()
+
+
 class Event:
     """Thread safe version of an async Event"""
 
@@ -547,39 +563,47 @@ def set(self) -> None:
             if not self._value:
                 self._value = True
 
-                for fut in self._waiters:
+                while self._waiters:
+                    fut = self._waiters.popleft()
+
                     if not fut.done():
                         if fut._loop == asyncio.get_running_loop():
                             if not fut.done():
                                 fut.set_result(True)
                         else:
 
-                            def s(w: asyncio.Future[Any]) -> None:
-                                if not w.done():
-                                    w.set_result(True)
+                            def s(w: asyncio.Future[Any], ev: threading.Event) -> None:
+                                try:
+                                    if not w.done():
+                                        w.set_result(True)
+                                finally:
+                                    ev.set()
 
                             if not fut.done():
-                                fut._loop.call_soon_threadsafe(s, fut)
+                                done = threading.Event()
+
+                                fut._loop.call_soon_threadsafe(s, fut, done)
+
+                                if not done.wait(120):
+                                    raise RuntimeError("Callback timeout")
 
     def clear(self) -> None:
         with self._lock:
             self._value = False
 
     async def wait(self, timeout: Optional[float] = None) -> bool:
-        if self._value:
-            return True
-
-        fut = create_sub_future()
         with self._lock:
+            if self._value:
+                return True
+
+            fut = create_sub_future()
             self._waiters.append(fut)
+
         try:
             await asyncio.wait_for(fut, timeout)
             return True
         except asyncio.TimeoutError:
             return False
-        finally:
-            with self._lock:
-                self._waiters.remove(fut)
 
 
 class Semaphore:
@@ -600,8 +624,8 @@ def __repr__(self) -> str:
             extra = f"{extra}, waiters:{len(self._waiters)}"
         return f"<{res[1:-1]} [{extra}]>"
 
-    def _wake_up_next(self) -> None:
-        with self._lock:
+    async def _wake_up_next(self) -> None:
+        async with async_lock(self._lock):
             while self._waiters:
                 waiter = self._waiters.popleft()
 
@@ -612,14 +636,23 @@ def _wake_up_next(self) -> None:
                     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)
+                            def s(w: asyncio.Future[Any], ev: threading.Event) -> None:
+                                try:
+                                    if w._loop.is_running() and not w.done():
+                                        w.set_result(True)
+                                finally:
+                                    ev.set()
 
                             if not waiter.done():
-                                waiter._loop.call_soon_threadsafe(s, waiter)
+                                done = threading.Event()
+
+                                waiter._loop.call_soon_threadsafe(s, waiter, done)
+
+                                if not done.wait(120):
+                                    raise RuntimeError("Callback timeout")
+
                         else:
-                            warnings.warn("Loop is not running.")
+                            raise RuntimeError("Loop is not running.")
 
     def locked(self) -> bool:
         with self._lock:
@@ -628,7 +661,7 @@ def locked(self) -> bool:
     async def acquire(self, timeout: Optional[float] = None) -> bool:
         while self._value <= 0:
             fut = create_sub_future()
-            with self._lock:
+            async with async_lock(self._lock):
                 self._waiters.append(fut)
             try:
                 await asyncio.wait_for(fut, timeout)
@@ -638,26 +671,26 @@ async def acquire(self, timeout: Optional[float] = None) -> bool:
                 if not fut.done():
                     fut.cancel()
                 if self._value > 0 and not fut.cancelled():
-                    self._wake_up_next()
+                    await self._wake_up_next()
 
                 raise
 
-        with self._lock:
+        async with async_lock(self._lock):
             self._value -= 1
 
         return True
 
-    def release(self) -> None:
+    async def release(self) -> None:
         self._value += 1
-        self._wake_up_next()
+        await self._wake_up_next()
 
     async def __aenter__(self) -> None:
         await self.acquire()
 
     async def __aexit__(
         self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
     ) -> None:
-        self.release()
+        await self.release()
 
 
 class BoundedSemaphore(Semaphore):
@@ -667,10 +700,10 @@ def __init__(self, value: int = 1) -> None:
         self._bound_value = value
         super().__init__(value)
 
-    def release(self) -> None:
+    async def release(self) -> None:
         if self._value >= self._bound_value:
             raise ValueError("BoundedSemaphore released too many times")
-        super().release()
+        await super().release()
 
 
 class Lock:
@@ -683,8 +716,8 @@ def __repr__(self) -> str:
     async def acquire(self, timeout: Optional[float] = None) -> bool:
         return await self._block.acquire(timeout)
 
-    def release(self) -> None:
-        self._block.release()
+    async def release(self) -> None:
+        await self._block.release()
 
     @property
     def locked(self) -> bool:
@@ -696,7 +729,7 @@ async def __aenter__(self) -> None:
     async def __aexit__(
         self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
     ) -> None:
-        self.release()
+        await self.release()
 
 
 class OldLock:
@@ -772,7 +805,7 @@ async def release(self) -> None:
                 if self._locked:
                     self._locked = False
                 else:
-                    warnings.warn(f"Lock is not acquired ({len(self._waiters) if self._waiters else 0} waiters).")
+                    raise RuntimeError(f"Lock is not acquired ({len(self._waiters) if self._waiters else 0} waiters).")
 
         await self._wake_up_next()