81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
|
"""In-memory WebSocket pub/sub with viewport filtering.
|
||
|
|
|
||
|
|
Zero extra deps. Ingest workers publish AIS/ADS-B points; only clients whose
|
||
|
|
current map bbox contains the point receive the payload. No Redis/Kafka.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
BBox = tuple[float, float, float, float] # minlon, minlat, maxlon, maxlat
|
||
|
|
|
||
|
|
|
||
|
|
def point_in_bbox(lon: float, lat: float, bbox: BBox | None) -> bool:
|
||
|
|
"""True if (lon, lat) sits inside an axis-aligned viewport."""
|
||
|
|
if bbox is None:
|
||
|
|
return False
|
||
|
|
minlon, minlat, maxlon, maxlat = bbox
|
||
|
|
return minlon <= lon <= maxlon and minlat <= lat <= maxlat
|
||
|
|
|
||
|
|
|
||
|
|
class ConnectionManager:
|
||
|
|
"""Maps Tailscale/browser clients → viewport bbox + per-client queue."""
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self._queues: dict[str, asyncio.Queue] = {}
|
||
|
|
self._viewports: dict[str, BBox] = {}
|
||
|
|
|
||
|
|
def register(self, client_id: str, maxsize: int = 256) -> asyncio.Queue:
|
||
|
|
q: asyncio.Queue = asyncio.Queue(maxsize=maxsize)
|
||
|
|
self._queues[client_id] = q
|
||
|
|
return q
|
||
|
|
|
||
|
|
def unregister(self, client_id: str) -> None:
|
||
|
|
self._queues.pop(client_id, None)
|
||
|
|
self._viewports.pop(client_id, None)
|
||
|
|
|
||
|
|
def set_viewport(self, client_id: str, bbox: BBox) -> None:
|
||
|
|
if client_id in self._queues:
|
||
|
|
self._viewports[client_id] = bbox
|
||
|
|
|
||
|
|
def viewport_of(self, client_id: str) -> BBox | None:
|
||
|
|
return self._viewports.get(client_id)
|
||
|
|
|
||
|
|
def has_clients(self) -> bool:
|
||
|
|
return bool(self._queues)
|
||
|
|
|
||
|
|
async def publish_point(
|
||
|
|
self,
|
||
|
|
kind: str,
|
||
|
|
payload: dict[str, Any],
|
||
|
|
*,
|
||
|
|
lat: float,
|
||
|
|
lon: float,
|
||
|
|
) -> int:
|
||
|
|
"""Enqueue `{type, payload}` for clients whose viewport contains the point.
|
||
|
|
|
||
|
|
Drops the oldest queued message if a client's buffer is full so a slow
|
||
|
|
tab cannot stall ingest. Returns the number of clients that got a copy.
|
||
|
|
"""
|
||
|
|
msg = {"type": kind, "payload": payload}
|
||
|
|
sent = 0
|
||
|
|
for client_id, queue in list(self._queues.items()):
|
||
|
|
if not point_in_bbox(lon, lat, self._viewports.get(client_id)):
|
||
|
|
continue
|
||
|
|
if queue.full():
|
||
|
|
try:
|
||
|
|
queue.get_nowait()
|
||
|
|
except asyncio.QueueEmpty:
|
||
|
|
pass
|
||
|
|
try:
|
||
|
|
queue.put_nowait(msg)
|
||
|
|
except asyncio.QueueFull:
|
||
|
|
continue
|
||
|
|
sent += 1
|
||
|
|
return sent
|
||
|
|
|
||
|
|
|
||
|
|
manager = ConnectionManager()
|