optimize cache and cancel algorithms

Author: d-biehl

package.json

@@ -119,6 +119,9 @@
           ],
           "escape": [
             "constant.character.escape.robotframework"
+          ],
+          "namespace": [
+            "entity.name.namespace.robotframework"
           ]
         }
       }
@@ -137,7 +140,7 @@
         "configuration": "./language-configuration.json"
       }
     ],
-    "#grammars": [
+    "grammars": [
       {
         "language": "robotframework",
         "scopeName": "source.robotframework",

robotcode/debugger/protocol.py

@@ -163,8 +163,8 @@ def done(f: asyncio.Future[Any]) -> None:
                 self._logger.exception(ex, exc_info=ex)
 
         for m in iterator:
-            future = asyncio.ensure_future(self.handle_message(m))
-            future.add_done_callback(done)
+            task = asyncio.create_task(self.handle_message(m))
+            task.add_done_callback(done)
 
     @_logger.call
     async def handle_message(self, message: ProtocolMessage) -> None:
@@ -262,11 +262,11 @@ async def handle_request(self, message: Request) -> None:
 
         try:
             if e is None or not callable(e.method):
-                result = asyncio.ensure_future(self.handle_unknown_command(message))
+                result = asyncio.create_task(self.handle_unknown_command(message))
             else:
                 params = self._convert_params(e.method, e.param_type, message.arguments)
 
-                result = asyncio.ensure_future(ensure_coroutine(e.method)(*params[0], **params[1]))
+                result = asyncio.create_task(ensure_coroutine(e.method)(*params[0], **params[1]))
 
             with self._received_request_lock:
                 self._received_request[message.seq] = result

robotcode/jsonrpc2/protocol.py

@@ -399,11 +399,11 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None:
         if isinstance(transport, asyncio.WriteTransport):
             self.write_transport = transport
 
-        asyncio.ensure_future(self.on_connection_made(self, transport))
+        asyncio.create_task(self.on_connection_made(self, transport))
 
     @_logger.call
     def connection_lost(self, exc: Optional[BaseException]) -> None:
-        asyncio.ensure_future(self.on_connection_lost(self, exc))
+        asyncio.create_task(self.on_connection_lost(self, exc))
 
     @_logger.call
     def eof_received(self) -> Optional[bool]:
@@ -508,8 +508,8 @@ def done(f: asyncio.Future[Any]) -> None:
                 self._logger.exception(ex, exc_info=ex)
 
         for m in iterator:
-            future = asyncio.ensure_future(self.handle_message(m))
-            future.add_done_callback(done)
+            task = asyncio.create_task(self.handle_message(m))
+            task.add_done_callback(done)
 
     @_logger.call
     async def handle_message(self, message: JsonRPCMessage) -> None:
@@ -698,7 +698,7 @@ async def handle_request(self, message: JsonRPCRequest) -> None:
         try:
             params = self._convert_params(e.method, e.param_type, message.params)
 
-            result = asyncio.ensure_future(ensure_coroutine(e.method)(*params[0], **params[1]))
+            result = asyncio.create_task(ensure_coroutine(e.method)(*params[0], **params[1]))
 
             with self._received_request_lock:
                 self._received_request[message.id] = result

robotcode/jsonrpc2/server.py

@@ -116,12 +116,12 @@ def start_stdio(self) -> None:
 
         protocol = self.create_protocol()
 
-        protocol.connection_made(transport)
-
         def run_io_nonblocking() -> None:
             self._stdio_stop_event = asyncio.Event()
 
             async def aio_readline(rfile: BinaryIO, protocol: asyncio.Protocol) -> None:
+                protocol.connection_made(transport)
+
                 while self._stdio_stop_event is not None and not self._stdio_stop_event.is_set() and not rfile.closed:
                     data = await self.loop.run_in_executor(None, cast(io.BufferedReader, rfile).read1, 1000)
                     protocol.data_received(data)

robotcode/language_server/common/parts/diagnostics.py

@@ -82,7 +82,7 @@ async def cancel() -> None:
                 except BaseException:
                     pass
 
-        asyncio.ensure_future(cancel())
+        asyncio.create_task(cancel())
 
 
 @dataclass

robotcode/language_server/common/text_document.py

@@ -38,6 +38,12 @@ class InvalidRangeError(Exception):
 _T = TypeVar("_T")
 
 
+class CacheEntry:
+    def __init__(self, data: Any = None) -> None:
+        self.data = data
+        self.lock: asyncio.Lock = asyncio.Lock()
+
+
 class TextDocument:
     def __init__(
         self,
@@ -73,7 +79,7 @@ def __init__(
 
         self._lines: Optional[List[str]] = None
 
-        self._cache: Dict[weakref.ref[Any], Any] = {}
+        self._cache: Dict[weakref.ref[Any], CacheEntry] = {}
         self._in_change_cache = False
 
         self._data: weakref.WeakKeyDictionary[Any, Any] = weakref.WeakKeyDictionary()
@@ -170,23 +176,6 @@ async def invalidate_data(self) -> None:
         async with self._lock:
             self._invalidate_data()
 
-    async def get_cache(
-        self,
-        entry: Union[Callable[[TextDocument], Awaitable[_T]], Callable[..., Awaitable[_T]]],
-        *args: Any,
-        **kwargs: Any,
-    ) -> _T:
-        if self._in_change_cache:
-            return await self._get_cache(entry, *args, **kwargs)
-
-        else:
-            self._in_change_cache = True
-            try:
-                async with self._lock:
-                    return await self._get_cache(entry, *args, **kwargs)
-            finally:
-                self._in_change_cache = False
-
     async def __remove_cache_entry_safe(self, ref: Any) -> None:
         async with self._lock:
             self._cache.pop(ref)
@@ -206,7 +195,7 @@ def __get_cache_reference(self, entry: Callable[..., Any]) -> weakref.ref[Any]:
 
         return reference
 
-    async def _get_cache(
+    async def get_cache(
         self,
         entry: Union[Callable[[TextDocument], Awaitable[_T]], Callable[..., Awaitable[_T]]],
         *args: Any,
@@ -216,42 +205,21 @@ async def _get_cache(
         reference = self.__get_cache_reference(entry)
 
         if reference not in self._cache:
-            self._cache[reference] = None
+            async with self._lock:
+                self._cache[reference] = CacheEntry()
 
-        if self._cache[reference] is None:
-            result = entry(self, *args, **kwargs)  # type: ignore
+        e = self._cache[reference]
 
-            if isinstance(result, Awaitable):
-                self._cache[reference] = await result
-            else:
-                self._cache[reference] = result
-
-        return cast("_T", self._cache[reference])
-
-    async def set_cache(
-        self,
-        entry: Union[Callable[[TextDocument], Awaitable[_T]], Callable[..., Awaitable[_T]]],
-        data: _T,
-    ) -> _T:
-        if self._in_change_cache:
-            return await self._set_cache(entry, data)
-
-        else:
-            self._in_change_cache = True
-            try:
-                async with self._lock:
-                    return await self._set_cache(entry, data)
-            finally:
-                self._in_change_cache = False
-
-    async def _set_cache(
-        self, entry: Union[Callable[[TextDocument], Awaitable[_T]], Callable[..., Awaitable[_T]]], data: _T
-    ) -> _T:
-        reference = self.__get_cache_reference(entry)
+        async with e.lock:
+            if e.data is None:
+                result = entry(self, *args, **kwargs)  # type: ignore
 
-        self._cache[reference] = data
+                if isinstance(result, Awaitable):
+                    e.data = await result
+                else:
+                    e.data = result
 
-        return data
+        return cast("_T", e.data)
 
     async def remove_cache_entry(
         self, entry: Union[Callable[[TextDocument], Awaitable[_T]], Callable[..., Awaitable[_T]]]