代码内容转移

backend_service/src/websocket_manager.py

@@ -0,0 +1,48 @@
+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()