"""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