Improved concurrency (#11)

theatre/src/theatre/server/connection.py

@@ -1,8 +1,8 @@
+import asyncio
 from asyncio import AbstractEventLoop
 from typing import List, Optional, Union
 
 from aiortc import (
-    InvalidStateError,
     RTCDataChannel,
     RTCPeerConnection,
     RTCSessionDescription,
@@ -61,13 +61,17 @@ def session(self) -> Session:
         return Session(description=self._peer_connection.localDescription.sdp)
 
     async def send(self, data: Union[str, bytes]) -> None:
-        for channel in self._channels:
-            try:
-                channel.send(data)
-                await channel.transport._data_channel_flush()
-                await channel.transport._transmit()
-            except (InvalidStateError, ConnectionError):
-                pass
+        async def send_to_channel(
+            data: Union[str, bytes], channel: RTCDataChannel
+        ) -> None:
+            channel.send(data)
+            await channel.transport._data_channel_flush()
+            await channel.transport._transmit()
+
+        tasks = [send_to_channel(data, channel) for channel in self._channels]
+
+        if len(tasks) > 0:
+            await asyncio.gather(*tasks, return_exceptions=True)
 
     async def close(self) -> None:
         await self._peer_connection.close()

theatre/src/theatre/server/room.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import random
 import uuid
@@ -90,11 +91,24 @@ def create_user(self) -> User:
         return User(id=user_id, avatar=avatar)
 
     async def broadcast(self, user_id: str, data: Union[str, bytes]) -> None:
-        for connected_user_id in self._connected_users:
-            if connected_user_id != user_id:
-                user_info = self._users.get(connected_user_id)
-                if user_info is not None:
-                    await user_info.connection.send(data)
+        async def send_to_user(
+            data: Union[str, bytes], user_info: UserInfo
+        ) -> None:
+            await user_info.connection.send(data)
+
+        user_infos = [
+            self._users.get(connected_user_id)
+            for connected_user_id in self._connected_users
+            if connected_user_id != user_id
+        ]
+        tasks = [
+            send_to_user(data, user_info)
+            for user_info in user_infos
+            if user_info is not None
+        ]
+
+        if len(tasks) > 0:
+            await asyncio.gather(*tasks, return_exceptions=True)
 
     async def handle_data(self, user_id: str, data: Union[str, bytes]) -> None:
         try: