222 lines
7.8 KiB
Python
222 lines
7.8 KiB
Python
|
J
|
"""Thin IQ-streaming agent.
|
||
|
|
|
||
|
|
Listens for control messages from the RIA Hub over a persistent WebSocket.
|
||
|
|
When the server sends ``start``, opens the SDR described in ``radio_config``,
|
||
|
|
loops over ``sdr.rx(buffer_size)``, and sends each buffer as raw
|
||
|
|
interleaved float32 bytes. ``stop`` closes the SDR; ``configure`` applies
|
||
|
|
parameter updates at the next capture boundary.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .hardware import heartbeat_payload
|
||
|
|
from .ws_client import WsClient
|
||
|
|
|
||
|
|
logger = logging.getLogger("ria_agent.streamer")
|
||
|
|
|
||
|
|
_DEFAULT_BUFFER_SIZE = 1024
|
||
|
|
|
||
|
|
|
||
|
|
class Streamer:
|
||
|
|
"""Main streamer loop.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
ws:
|
||
|
|
Connected :class:`WsClient`.
|
||
|
|
sdr_factory:
|
||
|
|
Callable ``(device, identifier) -> SDR``. Defaults to
|
||
|
|
:func:`ria_toolkit_oss.sdr.get_sdr_device`. Injectable for tests.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, ws: WsClient, sdr_factory=None) -> None:
|
||
|
|
self.ws = ws
|
||
|
|
self._sdr_factory = sdr_factory
|
||
|
|
self._app_id: str | None = None
|
||
|
|
self._sdr: Any = None
|
||
|
|
self._pending_config: dict = {}
|
||
|
|
self._capture_task: asyncio.Task | None = None
|
||
|
|
self._status = "idle"
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
# WsClient wiring
|
||
|
|
|
||
|
|
def build_heartbeat(self) -> dict:
|
||
|
|
return heartbeat_payload(status=self._status, app_id=self._app_id)
|
||
|
|
|
||
|
|
async def on_message(self, msg: dict) -> None:
|
||
|
|
t = msg.get("type")
|
||
|
|
if t == "start":
|
||
|
|
await self._handle_start(msg)
|
||
|
|
elif t == "stop":
|
||
|
|
await self._handle_stop(msg)
|
||
|
|
elif t == "configure":
|
||
|
|
self._pending_config.update(msg.get("radio_config") or {})
|
||
|
|
logger.debug("Queued configure: %s", self._pending_config)
|
||
|
|
else:
|
||
|
|
logger.warning("Unknown server message type: %r", t)
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
async def _handle_start(self, msg: dict) -> None:
|
||
|
|
if self._capture_task is not None and not self._capture_task.done():
|
||
|
|
logger.warning("start received while already streaming — ignoring")
|
||
|
|
return
|
||
|
|
|
||
|
|
self._app_id = msg.get("app_id")
|
||
|
|
radio_config = dict(msg.get("radio_config") or {})
|
||
|
|
device = radio_config.pop("device", None)
|
||
|
|
identifier = radio_config.pop("identifier", None)
|
||
|
|
buffer_size = int(radio_config.pop("buffer_size", _DEFAULT_BUFFER_SIZE))
|
||
|
|
if not device:
|
||
|
|
await self._send_error("start missing radio_config.device")
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
factory = self._sdr_factory or _default_sdr_factory
|
||
|
|
self._sdr = factory(device, identifier)
|
||
|
|
_apply_sdr_config(self._sdr, radio_config)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.exception("Failed to open SDR %r", device)
|
||
|
|
await self._send_error(f"SDR init failed: {exc}")
|
||
|
|
return
|
||
|
|
|
||
|
|
self._status = "streaming"
|
||
|
|
await self._send_status("streaming")
|
||
|
|
self._capture_task = asyncio.create_task(
|
||
|
|
self._capture_loop(buffer_size), name="ria-streamer-capture"
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _handle_stop(self, msg: dict) -> None:
|
||
|
|
if self._capture_task is not None:
|
||
|
|
self._capture_task.cancel()
|
||
|
|
try:
|
||
|
|
await self._capture_task
|
||
|
|
except (asyncio.CancelledError, Exception):
|
||
|
|
pass
|
||
|
|
self._capture_task = None
|
||
|
|
self._close_sdr()
|
||
|
|
self._app_id = None
|
||
|
|
self._status = "idle"
|
||
|
|
await self._send_status("idle")
|
||
|
|
|
||
|
|
async def _capture_loop(self, buffer_size: int) -> None:
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
try:
|
||
|
|
while True:
|
||
|
|
if self._pending_config:
|
||
|
|
cfg = self._pending_config
|
||
|
|
self._pending_config = {}
|
||
|
|
try:
|
||
|
|
_apply_sdr_config(self._sdr, cfg)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.warning("Applying configure failed: %s", exc)
|
||
|
|
|
||
|
|
try:
|
||
|
|
samples = await loop.run_in_executor(None, self._sdr.rx, buffer_size)
|
||
|
|
except Exception as exc:
|
||
|
|
from ria_toolkit_oss.sdr import SdrDisconnectedError
|
||
|
|
|
||
|
|
if isinstance(exc, SdrDisconnectedError):
|
||
|
|
logger.warning("SDR disconnected: %s", exc)
|
||
|
|
await self._send_error(f"SDR disconnected: {exc}")
|
||
|
|
else:
|
||
|
|
logger.exception("SDR rx error")
|
||
|
|
await self._send_error(f"SDR capture failed: {exc}")
|
||
|
|
break
|
||
|
|
|
||
|
|
payload = _samples_to_interleaved_float32(samples)
|
||
|
|
try:
|
||
|
|
await self.ws.send_bytes(payload)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.warning("Send failed: %s — ending capture", exc)
|
||
|
|
break
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
self._close_sdr()
|
||
|
|
|
||
|
|
def _close_sdr(self) -> None:
|
||
|
|
if self._sdr is None:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
self._sdr.close()
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
self._sdr = None
|
||
|
|
|
||
|
|
async def _send_status(self, status: str) -> None:
|
||
|
|
try:
|
||
|
|
await self.ws.send_json({"type": "status", "status": status, "app_id": self._app_id})
|
||
|
|
except Exception as exc:
|
||
|
|
logger.debug("Status send failed: %s", exc)
|
||
|
|
|
||
|
|
async def _send_error(self, message: str) -> None:
|
||
|
|
try:
|
||
|
|
await self.ws.send_json({"type": "error", "app_id": self._app_id, "message": message})
|
||
|
|
except Exception as exc:
|
||
|
|
logger.debug("Error-frame send failed: %s", exc)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Helpers
|
||
|
|
|
||
|
|
_CONFIG_ATTR_MAP = {
|
||
|
|
"sample_rate": ("sample_rate", "rx_sample_rate"),
|
||
|
|
"center_frequency": ("center_freq", "rx_center_frequency"),
|
||
|
|
"center_freq": ("center_freq", "rx_center_frequency"),
|
||
|
|
"gain": ("gain", "rx_gain"),
|
||
|
|
"bandwidth": ("bandwidth", "rx_bandwidth"),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _apply_sdr_config(sdr: Any, cfg: dict) -> None:
|
||
|
|
"""Apply a radio_config dict to an SDR, trying multiple attribute aliases."""
|
||
|
|
for key, value in cfg.items():
|
||
|
|
if value is None:
|
||
|
|
continue
|
||
|
|
attrs = _CONFIG_ATTR_MAP.get(key, (key,))
|
||
|
|
applied = False
|
||
|
|
for attr in attrs:
|
||
|
|
if hasattr(sdr, attr):
|
||
|
|
try:
|
||
|
|
setattr(sdr, attr, value)
|
||
|
|
applied = True
|
||
|
|
break
|
||
|
|
except Exception as exc:
|
||
|
|
logger.debug("setattr %s=%r failed: %s", attr, value, exc)
|
||
|
|
if not applied:
|
||
|
|
logger.debug("radio_config key %r ignored (no matching attr)", key)
|
||
|
|
|
||
|
|
|
||
|
|
def _samples_to_interleaved_float32(samples: Any) -> bytes:
|
||
|
|
"""Convert complex IQ samples (any numeric dtype) to interleaved float32 bytes."""
|
||
|
|
arr = np.asarray(samples)
|
||
|
|
if np.iscomplexobj(arr):
|
||
|
|
interleaved = np.empty(arr.size * 2, dtype=np.float32)
|
||
|
|
interleaved[0::2] = arr.real.astype(np.float32, copy=False).ravel()
|
||
|
|
interleaved[1::2] = arr.imag.astype(np.float32, copy=False).ravel()
|
||
|
|
return interleaved.tobytes()
|
||
|
|
return arr.astype(np.float32, copy=False).tobytes()
|
||
|
|
|
||
|
|
|
||
|
|
def _default_sdr_factory(device: str, identifier: str | None):
|
||
|
|
from ria_toolkit_oss.sdr import get_sdr_device
|
||
|
|
|
||
|
|
return get_sdr_device(device, ident=identifier)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Top-level entry
|
||
|
|
|
||
|
|
async def run_streamer(ws_url: str, token: str) -> None:
|
||
|
|
"""Connect to *ws_url* and run the streamer loop until cancelled."""
|
||
|
|
ws = WsClient(ws_url, token)
|
||
|
|
streamer = Streamer(ws)
|
||
|
|
await ws.run(streamer.on_message, streamer.build_heartbeat)
|