dora-teleop/dora_teleop/server.py

"""The console's web server, running inside this node.

The browser used to reach the node through `dora-websocket-transceiver`, which
meant a second node, a second port and a round trip through the dataflow for
every button press. This serves the page and holds the sockets directly, so the
console is one URL and the operator wires nothing.

**Nothing here runs on the dora thread.** aiohttp wants an asyncio loop and the
node's loop is a blocking `for event in node`, so the server owns a thread of
its own and the two sides meet at two places only: :meth:`Console.take`, which
the tick drains, and :meth:`Console.broadcast`, which hands a snapshot the other
way. Both are non-blocking by construction — a browser that stops reading must
never be able to stall the thread that stops the arm.
"""

import asyncio
import json
import queue
import threading
from pathlib import Path

from aiohttp import WSMsgType, web

from dora_teleop.log import log

INDEX = Path(__file__).parent / "web" / "index.html"

# How many console messages may pile up between two ticks before we start
# dropping them. The console sends a jog every 150 ms and the tick runs every
# 50, so reaching this at all means something is flooding us; an unbounded
# queue would answer that by growing until the node died.
INBOX_LIMIT = 256

# How many snapshots may wait on one socket. Telemetry is a snapshot, not a
# log, so a client too slow to keep up wants the newest one and not a backlog
# of the arm's recent past.
OUTBOX_LIMIT = 8


class _Client:
    """One connected browser, and the snapshots waiting to reach it."""

    def __init__(self) -> None:
        self.outbox: asyncio.Queue[str] = asyncio.Queue(maxsize=OUTBOX_LIMIT)

    def offer(self, text: str) -> None:
        """Queue one snapshot, discarding the stalest if this client is behind."""
        if self.outbox.full():
            self.outbox.get_nowait()
        self.outbox.put_nowait(text)


class Console:
    """The page, the sockets, and the thread they live on."""

    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        # Read every tick by the deadman: the last browser leaving has to stop
        # the arm, and a closed socket sends nothing to say so.
        self.clients = 0
        self._inbox: queue.Queue[dict] = queue.Queue(maxsize=INBOX_LIMIT)
        self._sockets: set[_Client] = set()
        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None
        self._ready = threading.Event()
        # Whatever stopped the server coming up, kept so `start` can raise it on
        # the caller's thread instead of leaving a traceback on stderr and a
        # node that believes it is serving a console.
        self._failure: BaseException | None = None

    def start(self) -> None:
        """Bring the server up, or raise saying why it could not.

        Refusing to start is the point. The port being taken used to leave a
        traceback on the server thread and a node that carried on announcing a
        console nobody could open — an operator typing the URL and getting
        nothing, with the log saying everything was fine.
        """
        self._thread = threading.Thread(target=self._run, name="teleop-console", daemon=True)
        self._thread.start()
        if not self._ready.wait(timeout=10.0):
            raise RuntimeError(f"the console did not come up on {self.host}:{self.port} within 10s")
        if self._failure is not None:
            raise self._failure
        log(f"console on http://{self.host}:{self.port}")

    def take(self) -> list[dict]:
        """Return every console message that arrived since the last call."""
        drained = []
        while True:
            try:
                drained.append(self._inbox.get_nowait())
            except queue.Empty:
                return drained

    def broadcast(self, payload: dict) -> None:
        """Hand one object to every connected browser."""
        if self._loop is None:
            return
        text = json.dumps(payload)
        self._loop.call_soon_threadsafe(self._offer, text)

    def stop(self) -> None:
        """Close the sockets and stop the thread."""
        if self._loop is None:
            return
        self._loop.call_soon_threadsafe(self._loop.stop)
        if self._thread is not None:
            self._thread.join(timeout=2.0)

    # --- everything below runs on the server thread ---

    def _run(self) -> None:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        self._loop = loop
        try:
            runner = loop.run_until_complete(self._listen())
        except BaseException as error:  # noqa: BLE001 — handed to `start` to raise
            self._failure = error
            self._loop = None
            loop.close()
            # Releases `start` at once rather than making it wait out its
            # timeout for a server that is never coming up.
            self._ready.set()
            return
        try:
            loop.run_forever()
        finally:
            loop.run_until_complete(runner.cleanup())
            loop.close()

    async def _listen(self) -> web.AppRunner:
        app = web.Application()
        app.router.add_get("/", self._page)
        app.router.add_get("/ws", self._socket)
        runner = web.AppRunner(app)
        await runner.setup()
        await web.TCPSite(runner, self.host, self.port).start()
        self._ready.set()
        return runner

    async def _page(self, request: web.Request) -> web.FileResponse:
        """Serve the console itself."""
        return web.FileResponse(INDEX)

    async def _socket(self, request: web.Request) -> web.WebSocketResponse:
        """Hold one browser's socket for as long as it stays."""
        socket = web.WebSocketResponse(heartbeat=20)
        await socket.prepare(request)
        client = _Client()
        self._sockets.add(client)
        self.clients = len(self._sockets)
        writer = asyncio.create_task(self._write(socket, client))
        try:
            async for received in socket:
                if received.type == WSMsgType.TEXT:
                    self._accept(received.data)
        finally:
            writer.cancel()
            self._sockets.discard(client)
            self.clients = len(self._sockets)
        return socket

    def _accept(self, text: str) -> None:
        """Park one message from a browser for the next tick to read."""
        try:
            payload = json.loads(text)
        except json.JSONDecodeError:
            log("ignoring a console message that was not JSON")
            return
        try:
            self._inbox.put_nowait(payload)
        except queue.Full:
            log("console inbox full, dropping a message")

    def _offer(self, text: str) -> None:
        for client in self._sockets:
            client.offer(text)

    async def _write(self, socket: web.WebSocketResponse, client: _Client) -> None:
        while True:
            await socket.send_str(await client.outbox.get())