use shared objects like loops db access etc

backend/src/utilities/packet_broadcaster.py

@@ -0,0 +1,140 @@
+# src/utilities/packet_broadcaster.py
+import asyncio
+import logging
+from typing import Dict, Any, List, Optional
+
+logger = logging.getLogger("packet_broadcaster")
+
+
+class PacketBroadcaster:
+    """
+    Simple in-process broadcaster:
+    - Maintains a set of subscriber asyncio.Queues (one per websocket connection).
+    - publish(msg) is run on the broadcaster's event loop.
+    - sync_publish(msg) is thread-safe and can be called from other threads / loops.
+
+    Note: create this on the FastAPI event loop (e.g. in startup) so that its lock and
+    operations run on that same loop.
+    """
+
+    def __init__(self, loop: asyncio.AbstractEventLoop, queue_maxsize: int = 1024):
+        self._loop = loop
+        self._queue_maxsize = queue_maxsize
+
+        # create lock and subscribers on the target loop to avoid cross-loop asyncio primitives
+        self._subscribers: List[asyncio.Queue] = []
+        # create lock bound to the same loop by scheduling its construction on that loop
+        self._lock: Optional[asyncio.Lock] = None
+        try:
+            # ensure lock is created on the given loop
+            def _make_lock():
+                self._lock = asyncio.Lock()
+
+            loop.call_soon_threadsafe(_make_lock)
+        except Exception:
+            # fallback — create in current loop if call_soon_threadsafe fails
+            self._lock = asyncio.Lock()
+
+        self._closed = False
+
+    async def subscribe(self) -> asyncio.Queue:
+        """
+        Create a subscriber queue and add it to the list.
+        Caller is expected to await on the returned queue to receive messages.
+        """
+        if self._closed:
+            raise RuntimeError("PacketBroadcaster is closed")
+
+        q: asyncio.Queue = asyncio.Queue(maxsize=self._queue_maxsize)
+        # wait until lock exists
+        while self._lock is None:
+            await asyncio.sleep(0)  # yield to event loop briefly
+
+        async with self._lock:
+            self._subscribers.append(q)
+        return q
+
+    async def unsubscribe(self, q: asyncio.Queue) -> None:
+        """
+        Remove a subscriber queue if present.
+        """
+        if self._lock is None:
+            return
+        async with self._lock:
+            try:
+                self._subscribers.remove(q)
+            except ValueError:
+                pass
+
+    async def publish(self, msg: Dict[str, Any]) -> None:
+        """
+        Publish msg to all subscribers (must be called on the broadcaster's loop).
+        We use put_nowait to avoid blocking. If a subscriber queue is full we drop
+        that subscriber's message to avoid backpressure.
+        """
+        if self._closed:
+            return
+
+        if self._lock is None:
+            # not initialized yet; nothing to do
+            return
+
+        async with self._lock:
+            subs = list(self._subscribers)
+
+        for q in subs:
+            try:
+                q.put_nowait(msg)
+            except asyncio.QueueFull:
+                # drop message for this subscriber
+                continue
+            except Exception as exc:
+                logger.exception("Unexpected error when publishing to subscriber: %s", exc)
+                # attempt to remove broken subscriber
+                try:
+                    async with self._lock:
+                        if q in self._subscribers:
+                            self._subscribers.remove(q)
+                except Exception:
+                    pass
+
+    def sync_publish(self, msg: Dict[str, Any]) -> None:
+        """
+        Thread-safe publish method: schedule publish(msg) on the broadcaster's loop.
+        Safe to call from other threads / event loops.
+
+        We schedule creation of the publish task on the broadcaster loop using
+        call_soon_threadsafe so that publish() runs on the correct loop.
+        """
+        if self._closed:
+            return
+
+        try:
+            # schedule the coroutine to run on the broadcaster loop
+            self._loop.call_soon_threadsafe(asyncio.create_task, self.publish(msg))
+        except Exception as exc:
+            # swallow errors but log for debugging
+            logger.exception("sync_publish failed to schedule publish: %s", exc)
+
+    async def close(self) -> None:
+        """
+        Close the broadcaster: mark closed, clear subscribers, and drain queues.
+        """
+        self._closed = True
+        if self._lock is None:
+            return
+        async with self._lock:
+            subs = list(self._subscribers)
+            self._subscribers.clear()
+
+        for q in subs:
+            try:
+                # optionally notify subscribers of closure by putting None (client must handle)
+                # q.put_nowait(None)
+                while not q.empty():
+                    try:
+                        q.get_nowait()
+                    except Exception:
+                        break
+            except Exception:
+                pass