malmert
8e6c759cb8
All checks were successful
Build and Deploy MITM Webserver / build (push) Successful in 1m39s
104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
"""In-process packet broadcaster for websocket subscribers."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
logger = logging.getLogger("packet_broadcaster")
|
|
|
|
_SHUTDOWN_SENTINEL: Dict[str, Any] = {"type": "__broadcaster_shutdown__"}
|
|
|
|
|
|
class PacketBroadcaster:
|
|
"""Manage subscriber queues and publish packet events."""
|
|
|
|
def __init__(self, loop: asyncio.AbstractEventLoop, queue_maxsize: int = 1024):
|
|
self._loop = loop
|
|
self._queue_maxsize = queue_maxsize
|
|
self._subscribers: List[asyncio.Queue] = []
|
|
self._lock: Optional[asyncio.Lock] = None
|
|
self._closed = False
|
|
|
|
try:
|
|
def _make_lock() -> None:
|
|
self._lock = asyncio.Lock()
|
|
|
|
loop.call_soon_threadsafe(_make_lock)
|
|
except Exception:
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def subscribe(self) -> asyncio.Queue:
|
|
"""Create and register a queue for one subscriber."""
|
|
if self._closed:
|
|
raise RuntimeError("PacketBroadcaster is closed")
|
|
|
|
queue: asyncio.Queue = asyncio.Queue(maxsize=self._queue_maxsize)
|
|
while self._lock is None:
|
|
await asyncio.sleep(0)
|
|
|
|
async with self._lock:
|
|
self._subscribers.append(queue)
|
|
|
|
return queue
|
|
|
|
async def unsubscribe(self, queue: asyncio.Queue) -> None:
|
|
"""Unregister a subscriber queue if it exists."""
|
|
if self._lock is None:
|
|
return
|
|
|
|
async with self._lock:
|
|
try:
|
|
self._subscribers.remove(queue)
|
|
except ValueError:
|
|
pass
|
|
|
|
async def publish(self, msg: Dict[str, Any]) -> None:
|
|
"""Publish one message to all current subscribers."""
|
|
if self._closed or self._lock is None:
|
|
return
|
|
|
|
async with self._lock:
|
|
subscribers = list(self._subscribers)
|
|
|
|
for queue in subscribers:
|
|
try:
|
|
queue.put_nowait(msg)
|
|
except asyncio.QueueFull:
|
|
continue
|
|
except Exception as exc:
|
|
logger.exception("Unexpected subscriber publish error: %s", exc)
|
|
try:
|
|
async with self._lock:
|
|
if queue in self._subscribers:
|
|
self._subscribers.remove(queue)
|
|
except Exception:
|
|
pass
|
|
|
|
def sync_publish(self, msg: Dict[str, Any]) -> None:
|
|
"""Thread-safe wrapper that schedules `publish` on the broadcaster loop."""
|
|
if self._closed:
|
|
return
|
|
|
|
try:
|
|
self._loop.call_soon_threadsafe(asyncio.create_task, self.publish(msg))
|
|
except Exception as exc:
|
|
logger.exception("sync_publish failed to schedule publish: %s", exc)
|
|
|
|
async def close(self) -> None:
|
|
"""Close the broadcaster and clear queued messages."""
|
|
self._closed = True
|
|
if self._lock is None:
|
|
return
|
|
|
|
async with self._lock:
|
|
subscribers = list(self._subscribers)
|
|
self._subscribers.clear()
|
|
|
|
for queue in subscribers:
|
|
try:
|
|
while not queue.empty():
|
|
queue.get_nowait()
|
|
queue.put_nowait(_SHUTDOWN_SENTINEL)
|
|
except Exception:
|
|
pass
|