Handle requests in a non-blocking fashion

fastapi_websocket_rpc/rpc_methods.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 import sys
+import time
 import typing
 import copy
 
@@ -97,3 +98,16 @@ async def get_response(self, call_id="") -> typing.Any:
 
     async def echo(self, text: str) -> str:
         return text
+
+    async def wait(self, seconds: float) -> str:
+        await asyncio.sleep(seconds)
+        return "hello world"
+
+    async def time_me(self, method: str, args: dict, count: int) -> float:
+        tasks = set()
+        start = time.monotonic()
+        for i in range(count):
+            tasks.add(asyncio.create_task(self.channel.call(method, args=args)))
+        await asyncio.gather(*tasks)
+        end = time.monotonic()
+        return end - start

fastapi_websocket_rpc/websocket_rpc_client.py

@@ -327,6 +327,7 @@ async def reader(self):
         """
         Read responses from socket worker
         """
+        handlers = set()
         try:
             while True:
                 raw_message = await self.ws.recv()
@@ -336,14 +337,20 @@ async def reader(self):
                     await self.close()
                     break
                 else:
-                    await self.channel.on_message(raw_message)
+                    handlers.add(asyncio.create_task(self.channel.on_message(raw_message)))
+                    to_remove = {h for h in handlers if h.done()}
+                    handlers -= to_remove
         # Graceful external termination options
         # task was canceled
         except asyncio.CancelledError:
             pass
         except Exception as err:
             logger.exception("RPC Reader task failed")
             raise
+        finally:
+            for handler in handlers:
+                handler.cancel()
+            await asyncio.gather(*handlers, return_exceptions=True)
 
     async def _keep_alive(self):
         try:

fastapi_websocket_rpc/websocket_rpc_endpoint.py

@@ -86,10 +86,13 @@ async def main_loop(self, websocket: WebSocket, client_id: str = None, **kwargs)
             channel.register_disconnect_handler(self._on_disconnect)
             # trigger connect handlers
             await channel.on_connect()
+            handlers = set()
             try:
                 while True:
                     data = await simple_websocket.recv()
-                    await channel.on_message(data)
+                    handlers.add(asyncio.create_task(channel.on_message(data)))
+                    to_remove = {h for h in handlers if h.done()}
+                    handlers -= to_remove
             except WebSocketDisconnect:
                 logger.info(
                     f"Client disconnected - {websocket.client.port} :: {channel.id}")
@@ -99,6 +102,10 @@ async def main_loop(self, websocket: WebSocket, client_id: str = None, **kwargs)
                 logger.info(
                     f"Client connection failed - {websocket.client.port} :: {channel.id}")
                 await self.handle_disconnect(websocket, channel)
+            finally:
+                for handler in handlers:
+                    handler.cancel()
+                await asyncio.gather(*handlers, return_exceptions=True)
         except:
             logger.exception(f"Failed to serve - {websocket.client.port}")
             self.manager.disconnect(websocket)

tests/advanced_rpc_test.py

@@ -65,3 +65,19 @@ async def test_recursive_rpc_calls(server):
         response = await client.other.get_response(call_id=remote_promise.result)
         # check the response we sent
         assert response.result['result'] == text
+
+
+@pytest.mark.asyncio
+async def test_simultaneous_rpc_calls_to_server(server):
+    async with WebSocketRpcClient(uri, RpcUtilityMethods(), default_response_timeout=10) as client:
+        start = time.monotonic()
+        await asyncio.gather(*(asyncio.create_task(client.other.wait(seconds=1)) for _ in range(3)))
+        end = time.monotonic()
+        assert end - start < 2
+
+
+@pytest.mark.asyncio
+async def test_simultaneous_rpc_calls_to_client(server):
+    async with WebSocketRpcClient(uri, RpcUtilityMethods(), default_response_timeout=10) as client:
+        time_taken = (await client.other.time_me(method="wait", args={"seconds": 1}, count=3)).result
+        assert time_taken < 2