path: root/packages/meshbay-node/src/meshbay_node/transport/webrtc/workers.py

"""Connection processes, seen from the node: more than one core for WebRTC.

aiortc is Python on one core. One process moves about 35 MB/s summed over all
its connections and, at that point, answers nothing else: the event loop that
carries the transfers is the one that keeps every other connection alive
(docs/MESHBAY_DESIGN.md §5.5). With `connection_processes` above 1 the node
starts that many `worker.py` processes and gives each new peer connection to the
least loaded. The session, with everything it knows, stays here.

What a process holds is the stack below the data channel. `RemotePeerConnection`
and `RemoteChannel` stand where aiortc's objects stand in a session and offer
the part of them a session uses, so no handler can tell which it was given.

Frames, both ways, on one loopback socket per process:

    4 bytes  length of what follows
    1 byte   type
    4 bytes  connection, a number this side draws
    ...      payload: JSON for control frames, raw bytes for SEND and MESSAGE

The socket is TCP on 127.0.0.1 because that is the one transport asyncio offers
identically on every platform the node runs on. A process proves it is the one
this node started by sending back 32 random bytes it was given on its stdin.
"""

import asyncio
import json
import logging
import secrets
import struct
import sys

from pyee.asyncio import AsyncIOEventEmitter

from meshbay_node.config import max_processes  # noqa: F401  (re-exported)
from meshbay_node.transport.webrtc.worker import (
    ANSWER,
    CHANNEL,
    CLOSE,
    CONFIG,
    FLOW,
    LOG,
    MESSAGE,
    OFFER,
    QUIT,
    SEND,
    STATE,
    STATS,
    pack,
    read_frame,
)

log = logging.getLogger("meshbay_node.transport.webrtc_server")

START_TIMEOUT = 30       # seconds for a process to import aiortc and call back
ANSWER_TIMEOUT = 30      # an ICE gather is bounded at 5 s; this is for a dead process


def _argv() -> list[str]:
    if getattr(sys, "frozen", False):
        return [sys.executable, "webrtc-worker"]
    return [sys.executable, "-m", "meshbay_node.transport.webrtc.worker"]


class _Description:
    def __init__(self, sdp: str, kind: str) -> None:
        self.sdp = sdp
        self.type = kind


class RemoteChannel(AsyncIOEventEmitter):
    """The data channel of a connection another process holds."""

    def __init__(self, peer: "RemotePeerConnection", label: str) -> None:
        super().__init__()
        self._peer = peer
        self.label = label
        self.readyState = "open"
        self._sent = 0
        self._flushed = 0

    @property
    def bufferedAmount(self) -> int:
        """Bytes sent here that the process has not yet put on the wire: what
        is in the socket, plus what its SCTP stack still queues."""
        return max(0, self._sent - self._flushed)

    def send(self, data: bytes) -> None:
        if self.readyState != "open":
            return
        self._sent += len(data)
        self._peer._worker.put(SEND, self._peer._id, data)


class RemotePeerConnection(AsyncIOEventEmitter):
    """The part of an RTCPeerConnection a session uses, for a connection that
    lives in a connection process."""

    def __init__(self, worker: "_Worker", conn: int, offer_sdp: str) -> None:
        super().__init__()
        self._worker = worker
        self._id = conn
        self.connectionState = "new"
        self.iceConnectionState = "new"
        self.remoteDescription = _Description(offer_sdp, "offer")
        self.localDescription: _Description | None = None
        # Read by `_get_remote_ip` and `sctp_stats.describe` in place of the
        # aiortc internals they reach into: as last reported by the process.
        self.remote_ip = ""
        self.sctp_line = ""
        self.channel: RemoteChannel | None = None
        self._answer: asyncio.Future = asyncio.get_running_loop().create_future()

    async def negotiate(self) -> str:
        """Send the offer to the process; the answer SDP, candidates included."""
        self._worker.put(OFFER, self._id, json.dumps(
            {"sdp": self.remoteDescription.sdp}).encode())
        try:
            body = await asyncio.wait_for(self._answer, ANSWER_TIMEOUT)
        except TimeoutError:
            self._gone()
            raise RuntimeError("Connection process did not answer") from None
        if "error" in body:
            self._gone()
            raise RuntimeError(body["error"])
        self.localDescription = _Description(body["sdp"], "answer")
        return body["sdp"]

    async def close(self) -> None:
        if self.connectionState == "closed":
            return
        self._worker.put(CLOSE, self._id)
        self._gone()

    def _gone(self) -> None:
        """This connection is over, whoever ended it. Said once."""
        self._worker.forget(self._id)
        if self.channel is not None:
            self.channel.readyState = "closed"
        if not self._answer.done():
            self._answer.set_result({"error": "Connection process stopped"})
        if self.connectionState != "closed":
            self.connectionState = self.iceConnectionState = "closed"
            self.emit("connectionstatechange")

    def _frame(self, kind: int, payload: bytes) -> None:
        if kind == MESSAGE:
            if self.channel is not None:
                self.channel.emit("message", payload)
        elif kind == FLOW:
            if self.channel is not None:
                self.channel._flushed = struct.unpack(">Q", payload)[0]
        elif kind == STATS:
            self.sctp_line = payload.decode()
        elif kind == CHANNEL:
            self.channel = RemoteChannel(self, json.loads(payload)["label"])
            self.emit("datachannel", self.channel)
        elif kind == ANSWER:
            if not self._answer.done():
                self._answer.set_result(json.loads(payload))
        elif kind == STATE:
            body = json.loads(payload)
            self.remote_ip = body.get("ip") or self.remote_ip
            if self.channel is not None and body.get("channel") == "closed":
                self.channel.readyState = "closed"
            if body["connection"] == "closed":
                self._gone()
                return
            ice, self.iceConnectionState = self.iceConnectionState, body["ice"]
            if ice != body["ice"]:
                self.emit("iceconnectionstatechange")
            was, self.connectionState = self.connectionState, body["connection"]
            if was != body["connection"]:
                if body["connection"] == "failed":
                    self._worker.forget(self._id)
                self.emit("connectionstatechange")


class _Worker:
    def __init__(self, number: int) -> None:
        self.number = number
        self.process: asyncio.subprocess.Process | None = None
        self.writer: asyncio.StreamWriter | None = None
        self.conns: dict[int, RemotePeerConnection] = {}
        self.draining = False
        self.alive = False
        self.on_empty = None
        self._task: asyncio.Task | None = None

    def put(self, kind: int, conn: int, payload: bytes = b"") -> None:
        if self.alive and self.writer is not None and not self.writer.is_closing():
            self.writer.write(pack(kind, conn, payload))

    def forget(self, conn: int) -> None:
        if self.conns.pop(conn, None) is not None and not self.conns \
                and self.draining and self.on_empty is not None:
            self.on_empty(self)

    async def pump(self, reader: asyncio.StreamReader) -> None:
        try:
            while True:
                kind, conn, payload = await read_frame(reader)
                if kind == LOG:
                    body = json.loads(payload)
                    logging.getLogger(body["name"]).log(
                        body["level"], "[conn %d] %s", self.number, body["text"])
                    continue
                peer = self.conns.get(conn)
                if peer is not None:
                    try:
                        peer._frame(kind, payload)
                    except Exception:
                        log.exception("connection process %d: frame %d", self.number, kind)
        except (asyncio.IncompleteReadError, ConnectionError, ValueError):
            pass
        finally:
            self.stopped()
            # Reaped here, whichever way it went: asked to leave, killed, or
            # fallen over. A process nobody waits for stays a zombie.
            try:
                await asyncio.wait_for(self.process.wait(), 5)
            except TimeoutError:
                self.process.kill()

    def stopped(self) -> None:
        """The process is gone: every connection it held is gone with it."""
        was, self.alive = self.alive, False
        if was and not self.draining:
            log.warning("connection process %d stopped with %d connections",
                        self.number, len(self.conns))
        for peer in list(self.conns.values()):
            peer._gone()
        self.conns.clear()
        if self.writer is not None:
            self.writer.close()
        if self.process is not None and self.process.returncode is None:
            try:
                self.process.terminate()
            except ProcessLookupError:
                pass


class WorkerPool:
    """The connection processes of one node."""

    def __init__(self, count: int = 1, stun: list[str] | None = None,
                 ice_interfaces: list[str] | None = None) -> None:
        self.count = max(1, int(count))
        self._stun = list(stun or [])
        self._ice = list(ice_interfaces or [])
        self._workers: list[_Worker] = []
        self._leaving: list[_Worker] = []      # asked to go, not yet reaped
        self._server: asyncio.base_events.Server | None = None
        self._expected: dict[bytes, asyncio.Future] = {}
        self._next_conn = 0
        self._next_worker = 0
        self._lock = asyncio.Lock()

    @property
    def enabled(self) -> bool:
        return self.count > 1

    def processes(self) -> int:
        return sum(1 for w in self._workers if w.alive)

    async def _accept(self, reader, writer) -> None:
        try:
            token = await asyncio.wait_for(reader.readexactly(32), 10)
        except (TimeoutError, asyncio.IncompleteReadError, ConnectionError):
            writer.close()
            return
        waiting = self._expected.pop(token, None)
        if waiting is None or waiting.done():
            writer.close()
            return
        waiting.set_result((reader, writer))

    async def _start_one(self) -> _Worker:
        if self._server is None:
            self._server = await asyncio.start_server(self._accept, "127.0.0.1", 0)
        port = self._server.sockets[0].getsockname()[1]
        token = secrets.token_bytes(32)
        waiting = asyncio.get_running_loop().create_future()
        self._expected[token] = waiting
        self._next_worker += 1
        worker = _Worker(self._next_worker)
        worker.on_empty = self._retire
        worker.process = await asyncio.create_subprocess_exec(
            *_argv(), stdin=asyncio.subprocess.PIPE)
        hello = {"port": port, "token": token.hex(), "stun": self._stun,
                 "ice_interfaces": self._ice,
                 "level": logging.getLogger().getEffectiveLevel()}
        worker.process.stdin.write(json.dumps(hello).encode() + b"\n")
        try:
            reader, worker.writer = await asyncio.wait_for(waiting, START_TIMEOUT)
        except TimeoutError:
            self._expected.pop(token, None)
            worker.stopped()
            raise RuntimeError("Connection process did not start") from None
        worker.alive = True
        worker._task = asyncio.ensure_future(worker.pump(reader))
        log.info("connection process %d started (pid %d)",
                 worker.number, worker.process.pid)
        return worker

    def _retire(self, worker: _Worker) -> None:
        if worker in self._workers:
            self._workers.remove(worker)
            self._leaving.append(worker)
        if worker.alive:
            log.info("connection process %d retired", worker.number)
            worker.put(QUIT, 0)
        worker.draining = True

    async def ensure(self) -> None:
        """Bring the number of processes taking connections to `count`."""
        async with self._lock:
            self._workers = [w for w in self._workers if w.alive]
            self._leaving = [w for w in self._leaving
                             if w._task is not None and not w._task.done()]
            taking = [w for w in self._workers if not w.draining]
            want = self.count if self.enabled else 0
            for worker in taking[want:]:
                worker.draining = True
                if not worker.conns:
                    self._retire(worker)
            missing = want - len(taking)
            if missing > 0:
                started = await asyncio.gather(
                    *(self._start_one() for _ in range(missing)),
                    return_exceptions=True)
                for worker in started:
                    if isinstance(worker, _Worker):
                        self._workers.append(worker)
                    else:
                        log.error("connection process failed to start: %s", worker)

    async def resize(self, count: int) -> None:
        """Set the number of processes. Connections already open stay where
        they are: a process asked to leave does so when its last one closes."""
        self.count = max(1, int(count))
        await self.ensure()

    def configure(self, *, stun: list[str] | None = None,
                  ice_interfaces: list[str] | None = None) -> None:
        body: dict = {}
        if stun is not None:
            self._stun = body["stun"] = list(stun)
        if ice_interfaces is not None:
            self._ice = body["ice_interfaces"] = list(ice_interfaces)
        for worker in self._workers:
            worker.put(CONFIG, 0, json.dumps(body).encode())

    async def open(self, offer_sdp: str) -> RemotePeerConnection | None:
        """A connection in the least loaded process, not yet negotiated.

        None when no process can take it: the caller then holds the connection
        itself, which is what a node with one process does anyway.
        """
        taking = [w for w in self._workers if w.alive and not w.draining]
        if len(taking) < self.count:
            await self.ensure()
            taking = [w for w in self._workers if w.alive and not w.draining]
        if not taking:
            return None
        worker = min(taking, key=lambda w: len(w.conns))
        self._next_conn += 1
        peer = RemotePeerConnection(worker, self._next_conn, offer_sdp)
        worker.conns[peer._id] = peer
        return peer

    def loads(self) -> list[int]:
        return [len(w.conns) for w in self._workers if w.alive]

    async def close(self) -> None:
        workers = self._workers + self._leaving
        self._workers, self._leaving = [], []
        for worker in workers:
            worker.draining = True
            worker.put(QUIT, 0)
        for worker in workers:
            if worker.process is not None:
                try:
                    await asyncio.wait_for(worker.process.wait(), 5)
                except TimeoutError:
                    worker.process.kill()
            worker.stopped()
            if worker._task is not None:
                await asyncio.gather(worker._task, return_exceptions=True)
        if self._server is not None:
            self._server.close()
            self._server = None