49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
import asyncio
|
|
from typing import List
|
|
from fastapi import WebSocket
|
|
import logging
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: List[WebSocket] = []
|
|
self.loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
def set_loop(self, loop: asyncio.AbstractEventLoop):
|
|
"""Sets the asyncio event loop."""
|
|
self.loop = loop
|
|
|
|
async def connect(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.active_connections.append(websocket)
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
self.active_connections.remove(websocket)
|
|
|
|
def broadcast(self, message: str):
|
|
"""
|
|
Thread-safely broadcasts a message to all active WebSocket connections.
|
|
This method is designed to be called from a different thread (e.g., a ROS2 callback).
|
|
"""
|
|
if not self.loop:
|
|
logging.error("Event loop not set in ConnectionManager. Cannot broadcast.")
|
|
return
|
|
|
|
# Schedule the coroutine to be executed in the event loop
|
|
self.loop.call_soon_threadsafe(self._broadcast_in_loop, message)
|
|
|
|
def _broadcast_in_loop(self, message: str):
|
|
"""
|
|
Helper to run the broadcast coroutine in the correct event loop.
|
|
"""
|
|
asyncio.ensure_future(self._broadcast_async(message), loop=self.loop)
|
|
|
|
async def _broadcast_async(self, message: str):
|
|
"""
|
|
The actual async method that sends messages.
|
|
"""
|
|
tasks = [connection.send_text(message) for connection in self.active_connections]
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# Create a single instance of the manager to be used across the application
|
|
websocket_manager = ConnectionManager()
|