From d53fc46ffc515fda48fef5532c398d44b4f7207a Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Wed, 2 Mar 2022 23:51:12 -0800
Subject: [PATCH] All multiple event steam connections.

We don't support multi-client yet, but this is useful when debugging the
react UI in a browser instead of in the Liberation UI.
---
 game/server/eventstream/routes.py | 42 +++++++++++++++++++++++++------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/game/server/eventstream/routes.py b/game/server/eventstream/routes.py
index 2cbbfebb..e1210bfd 100644
--- a/game/server/eventstream/routes.py
+++ b/game/server/eventstream/routes.py
@@ -1,6 +1,7 @@
-from fastapi import APIRouter
+from asyncio import wait
+
+from fastapi import APIRouter, WebSocket
 from fastapi.encoders import jsonable_encoder
-from starlette.websockets import WebSocket
 
 from .eventstream import EventStream
 from .models import GameUpdateEventsJs
@@ -9,17 +10,42 @@ from .. import GameContext
 router: APIRouter = APIRouter()
 
 
+class ConnectionManager:
+    def __init__(self) -> None:
+        self.active_connections: list[WebSocket] = []
+
+    async def shutdown(self) -> None:
+        futures = []
+        for connection in self.active_connections:
+            futures.append(connection.close())
+        await wait(futures)
+
+    async def connect(self, websocket: WebSocket) -> None:
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket) -> None:
+        self.active_connections.remove(websocket)
+
+    async def broadcast(self, events: GameUpdateEventsJs) -> None:
+        futures = []
+        for connection in self.active_connections:
+            futures.append(connection.send_json(jsonable_encoder(events)))
+        await wait(futures)
+
+
+manager = ConnectionManager()
+
+
 @router.websocket("/eventstream")
 async def event_stream(websocket: WebSocket) -> None:
-    await websocket.accept()
+    await manager.connect(websocket)
     while True:
         if not (events := await EventStream.get()).empty:
             if events.shutting_down:
-                await websocket.close()
+                await manager.shutdown()
                 return
 
-            await websocket.send_json(
-                jsonable_encoder(
-                    GameUpdateEventsJs.from_events(events, GameContext.get())
-                )
+            await manager.broadcast(
+                GameUpdateEventsJs.from_events(events, GameContext.get())
             )