add async interface

Author: davschul

pylsp_jsonrpc/endpoint.py

@@ -4,7 +4,9 @@
 import logging
 import uuid
 import sys
+import asyncio
 
+from threading import Thread
 from concurrent import futures
 from .exceptions import (JsonRpcException, JsonRpcRequestCancelled,
                          JsonRpcInternalError, JsonRpcMethodNotFound)
@@ -14,6 +16,34 @@
 CANCEL_METHOD = '$/cancelRequest'
 
 
+async def run_as_daemon(func, *args):
+    future = futures.Future()
+    future.set_running_or_notify_cancel()
+
+    # A bug in python 3.7 makes it a bad idea to set a BaseException
+    # in a wrapped future (see except statement in asyncio.Task.__wakeup)
+    # Instead, we'll wrap base exceptions into exceptions and unwrap them
+    # on the other side of the call.
+    class BaseExceptionWrapper(Exception):
+        pass
+
+    def daemon():
+        try:
+            result = func(*args)
+        except Exception as e:
+            future.set_exception(e)
+        except BaseException as e:
+            future.set_exception(BaseExceptionWrapper(e))
+        else:
+            future.set_result(result)
+
+    Thread(target=daemon, daemon=True).start()
+    try:
+        return await asyncio.wrap_future(future)
+    except BaseExceptionWrapper as exc:
+        raise exc.args[0]
+
+
 class Endpoint:
 
     def __init__(self, dispatcher, consumer, id_generator=lambda: str(uuid.uuid4()), max_workers=5):
@@ -35,6 +65,19 @@ def __init__(self, dispatcher, consumer, id_generator=lambda: str(uuid.uuid4()),
         self._client_request_futures = {}
         self._server_request_futures = {}
         self._executor_service = futures.ThreadPoolExecutor(max_workers=max_workers)
+        self._cancelledRequests = set()
+
+    def init_async(self):
+        log.warning("init async")
+        self._messageQueue = asyncio.Queue()
+
+    async def consume_task(self):
+        log.warning("starting task")
+        while True:
+            message = await self._messageQueue.get()
+            await run_as_daemon(self.consume, message)
+            log.warning("got message in task")
+            self._messageQueue.task_done()
 
     def shutdown(self):
         self._executor_service.shutdown()
@@ -94,7 +137,15 @@ def callback(future):
                 future.set_exception(JsonRpcRequestCancelled())
         return callback
 
+    async def consume_async(self, message):
+        log.warning("got message put in queue")
+        if message['method'] == CANCEL_METHOD:
+            self._cancelledRequests.add(message.get('params')['id'])
+        await self._messageQueue.put(message)
+
+
     def consume(self, message):
+        log.warning("consume message")
         """Consume a JSON RPC message from the client.
 
         Args:
@@ -182,6 +233,9 @@ def _handle_request(self, msg_id, method, params):
         except KeyError as e:
             raise JsonRpcMethodNotFound.of(method) from e
 
+        if msg_id in self._cancelledRequests:
+            raise JsonRpcRequestCancelled()
+
         handler_result = handler(params)
 
         if callable(handler_result):

pylsp_jsonrpc/streams.py

@@ -3,6 +3,10 @@
 
 import logging
 import threading
+import asyncio
+import sys
+from threading import Thread
+from concurrent.futures import Future
 
 try:
     import ujson as json
@@ -12,6 +16,34 @@
 log = logging.getLogger(__name__)
 
 
+async def run_as_daemon(func, *args):
+    future = Future()
+    future.set_running_or_notify_cancel()
+
+    # A bug in python 3.7 makes it a bad idea to set a BaseException
+    # in a wrapped future (see except statement in asyncio.Task.__wakeup)
+    # Instead, we'll wrap base exceptions into exceptions and unwrap them
+    # on the other side of the call.
+    class BaseExceptionWrapper(Exception):
+        pass
+
+    def daemon():
+        try:
+            result = func(*args)
+        except Exception as e:
+            future.set_exception(e)
+        except BaseException as e:
+            future.set_exception(BaseExceptionWrapper(e))
+        else:
+            future.set_result(result)
+
+    Thread(target=daemon, daemon=True).start()
+    try:
+        return await asyncio.wrap_future(future)
+    except BaseExceptionWrapper as exc:
+        raise exc.args[0]
+
+
 class JsonRpcStreamReader:
     def __init__(self, rfile):
         self._rfile = rfile
@@ -25,6 +57,7 @@ def listen(self, message_consumer):
         Args:
             message_consumer (fn): function that is passed each message as it is read off the socket.
         """
+
         while not self._rfile.closed:
             try:
                 request_str = self._read_message()
@@ -65,6 +98,53 @@ def _read_message(self):
         # Grab the body
         return self._rfile.read(content_length)
 
+    async def listen_async(self, message_consumer):
+        """Blocking call to listen for messages on the rfile.
+
+        Args:
+            message_consumer (fn): function that is passed each message as it is read off the socket.
+        """
+
+        while not self._rfile.closed:
+            try:
+                request_str = await self._read_message_async()
+            except ValueError:
+                if self._rfile.closed:
+                    return
+                log.exception("Failed to read from rfile")
+
+            if request_str is None:
+                break
+
+            try:
+                await message_consumer(json.loads(request_str.decode('utf-8')))
+            except ValueError:
+                log.exception("Failed to parse JSON message %s", request_str)
+                continue
+
+    async def _read_message_async(self):
+        """Reads the contents of a message.
+
+        Returns:
+            body of message if parsable else None
+        """
+        line = await run_as_daemon(self._rfile.readline)
+
+        if not line:
+            return None
+
+        content_length = self._content_length(line)
+
+        # Blindly consume all header lines
+        while line and line.strip():
+            line = await run_as_daemon(self._rfile.readline)
+
+        if not line:
+            return None
+
+        # Grab the body
+        return await run_as_daemon(self._rfile.read, content_length)
+
     @staticmethod
     def _content_length(line):
         """Extract the content length from an input line."""