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.
This commit is contained in:
Dan Albert 2022-03-02 23:51:12 -08:00
parent 781f8fb0e8
commit d53fc46ffc

View File

@ -1,6 +1,7 @@
from fastapi import APIRouter from asyncio import wait
from fastapi import APIRouter, WebSocket
from fastapi.encoders import jsonable_encoder from fastapi.encoders import jsonable_encoder
from starlette.websockets import WebSocket
from .eventstream import EventStream from .eventstream import EventStream
from .models import GameUpdateEventsJs from .models import GameUpdateEventsJs
@ -9,17 +10,42 @@ from .. import GameContext
router: APIRouter = APIRouter() 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") @router.websocket("/eventstream")
async def event_stream(websocket: WebSocket) -> None: async def event_stream(websocket: WebSocket) -> None:
await websocket.accept() await manager.connect(websocket)
while True: while True:
if not (events := await EventStream.get()).empty: if not (events := await EventStream.get()).empty:
if events.shutting_down: if events.shutting_down:
await websocket.close() await manager.shutdown()
return return
await websocket.send_json( await manager.broadcast(
jsonable_encoder( GameUpdateEventsJs.from_events(events, GameContext.get())
GameUpdateEventsJs.from_events(events, GameContext.get())
)
) )