diff options
Diffstat (limited to 'packages/meshbay-node/src/meshbay_node/transport/webrtc')
4 files changed, 643 insertions, 0 deletions
diff --git a/packages/meshbay-node/src/meshbay_node/transport/webrtc/channel.py b/packages/meshbay-node/src/meshbay_node/transport/webrtc/channel.py index 31cc83c..4bb4e08 100644 --- a/packages/meshbay-node/src/meshbay_node/transport/webrtc/channel.py +++ b/packages/meshbay-node/src/meshbay_node/transport/webrtc/channel.py @@ -82,6 +82,9 @@ class _DataChannelBuffer: def _get_remote_ip(pc: RTCPeerConnection) -> str: """Best-effort extraction of the remote peer IP from the ICE transport.""" + reported = getattr(pc, "remote_ip", None) + if reported is not None: # held by a connection process (workers.py) + return reported try: dtls = pc.sctp and pc.sctp.transport ice = dtls and dtls.transport diff --git a/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_stats.py b/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_stats.py index 338ded9..cc2eb1e 100644 --- a/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_stats.py +++ b/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_stats.py @@ -33,6 +33,9 @@ RTCSctpTransport._t3_expired = _counting_t3_expired def describe(pc) -> str: """The association of `pc` as a log suffix, or nothing when it has none.""" + reported = getattr(pc, "sctp_line", None) + if reported is not None: # held by a connection process (workers.py) + return reported sctp = getattr(pc, "sctp", None) if sctp is None: return "" diff --git a/packages/meshbay-node/src/meshbay_node/transport/webrtc/worker.py b/packages/meshbay-node/src/meshbay_node/transport/webrtc/worker.py new file mode 100644 index 0000000..072c849 --- /dev/null +++ b/packages/meshbay-node/src/meshbay_node/transport/webrtc/worker.py @@ -0,0 +1,245 @@ +"""A connection process: peer connections held for the node, and nothing else. + +The WebRTC stack (ICE, DTLS, SCTP) is Python and saturates one core at about +35 MB/s summed over every connection (docs/MESHBAY_DESIGN.md §5.5). This process +runs that stack and only that. It holds no node key, no group key and no index: +what crosses it is what crosses the wire, bytes of the data channel in both +directions, and the node proper decides what they mean. `workers.py` is the +other end and documents the frames. + +Started by the node with one JSON line on stdin: the loopback port to call +back, the token that proves it is the process the node started, and the ICE +settings. It ends when that socket does, so it cannot outlive its node. +""" + +import asyncio +import json +import logging +import struct +import sys + +HEADER = struct.Struct(">IBI") # length of what follows, frame type, connection +MAX_FRAME = 64 * 1024 * 1024 + +# node -> process +OFFER, SEND, CLOSE, CONFIG, QUIT = 1, 2, 3, 4, 5 +# process -> node +ANSWER, CHANNEL, MESSAGE, STATE, FLOW, STATS, LOG = 16, 17, 18, 19, 20, 21, 22 + +# How often the node learns what a channel has put on the wire. The download +# path waits for room in steps of 50 ms (files.py), so anything well under that +# is as good as reading `bufferedAmount` in place. +FLOW_INTERVAL = 0.02 +STATS_EVERY = 50 # flow ticks between two SCTP summaries + + +def pack(kind: int, conn: int, payload: bytes = b"") -> bytes: + return HEADER.pack(len(payload) + 5, kind, conn) + payload + + +async def read_frame(reader: asyncio.StreamReader) -> tuple[int, int, bytes]: + length, kind, conn = HEADER.unpack(await reader.readexactly(HEADER.size)) + if not 5 <= length <= MAX_FRAME: + raise ValueError(f"frame of {length} bytes") + return kind, conn, await reader.readexactly(length - 5) + + +class _Forward(logging.Handler): + """Hands every record to the node, which owns the log file.""" + + def __init__(self, write) -> None: + super().__init__() + self._write = write + + def emit(self, record: logging.LogRecord) -> None: + try: + body = {"name": record.name, "level": record.levelno, + "text": record.getMessage()} + self._write(pack(LOG, 0, json.dumps(body).encode())) + except Exception: + pass + + +class _Conn: + def __init__(self, pc) -> None: + self.pc = pc + self.channel = None + self.taken = 0 # bytes of SEND frames handed to the channel + self.reported = -1 + self.stats = "" + + +class Worker: + def __init__(self, reader, writer) -> None: + self._reader = reader + self._writer = writer + self._conns: dict[int, _Conn] = {} + self._stun: list[str] = [] + self._tasks: set[asyncio.Task] = set() + + def _put(self, kind: int, conn: int, payload: bytes = b"") -> None: + if not self._writer.is_closing(): + self._writer.write(pack(kind, conn, payload)) + + def _put_json(self, kind: int, conn: int, body: dict) -> None: + self._put(kind, conn, json.dumps(body).encode()) + + def _spawn(self, coro) -> None: + task = asyncio.ensure_future(coro) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + + def configure(self, cfg: dict) -> None: + from meshbay_node.transport import ice_filter, stun_multi + if "ice_interfaces" in cfg: + ice_filter.install(cfg["ice_interfaces"] or None) + if "stun" in cfg: + self._stun = list(cfg["stun"] or []) + stun_multi.install(self._stun or None) + + async def run(self) -> None: + flow = asyncio.ensure_future(self._report_flow()) + try: + while True: + kind, conn, payload = await read_frame(self._reader) + if kind == SEND: + c = self._conns.get(conn) + if c is not None and c.channel is not None \ + and c.channel.readyState == "open": + c.channel.send(payload) + if c is not None: + # Counted whether or not it went out: the node's view of + # what is buffered must come back to zero on a channel + # that has closed, or its sender waits for room forever. + c.taken += len(payload) + elif kind == OFFER: + self._spawn(self._offer(conn, json.loads(payload))) + elif kind == CLOSE: + self._spawn(self._close(conn)) + elif kind == CONFIG: + self.configure(json.loads(payload)) + elif kind == QUIT: + break + except (asyncio.IncompleteReadError, ConnectionError): + pass + finally: + flow.cancel() + for conn in list(self._conns): + await self._close(conn) + + async def _offer(self, conn: int, body: dict) -> None: + from aiortc import ( + RTCConfiguration, + RTCIceServer, + RTCPeerConnection, + RTCSessionDescription, + ) + + from meshbay_node.transport.webrtc.channel import _get_remote_ip + try: + config = RTCConfiguration( + iceServers=[RTCIceServer(urls=s) for s in self._stun]) + pc = RTCPeerConnection(configuration=config) + c = self._conns[conn] = _Conn(pc) + + def state() -> None: + self._put_json(STATE, conn, { + "connection": pc.connectionState, "ice": pc.iceConnectionState, + "ip": _get_remote_ip(pc), + "channel": c.channel.readyState if c.channel else ""}) + + @pc.on("datachannel") + def on_datachannel(channel) -> None: + c.channel = channel + self._put_json(CHANNEL, conn, {"label": channel.label}) + + @channel.on("message") + def on_message(message) -> None: + if isinstance(message, str): + message = message.encode() + self._put(MESSAGE, conn, message) + + channel.on("close", state) + + @pc.on("connectionstatechange") + async def on_state() -> None: + state() + if pc.connectionState in ("failed", "closed"): + if self._conns.pop(conn, None) is not None \ + and pc.connectionState == "failed": + await pc.close() + + pc.on("iceconnectionstatechange", state) + await pc.setRemoteDescription( + RTCSessionDescription(sdp=body["sdp"], type="offer")) + await pc.setLocalDescription(await pc.createAnswer()) + self._put_json(ANSWER, conn, {"sdp": pc.localDescription.sdp}) + except Exception as e: + gone = self._conns.pop(conn, None) + self._put_json(ANSWER, conn, {"error": str(e) or type(e).__name__}) + if gone is not None: + try: + await gone.pc.close() + except Exception: + pass + + async def _close(self, conn: int) -> None: + c = self._conns.pop(conn, None) + if c is None: + return + try: + await c.pc.close() + except Exception: + pass + self._put_json(STATE, conn, {"connection": "closed", "ice": "closed", + "ip": "", "channel": "closed"}) + + async def _report_flow(self) -> None: + from meshbay_node.transport.webrtc import sctp_stats + tick = 0 + while True: + await asyncio.sleep(FLOW_INTERVAL) + tick += 1 + for conn, c in list(self._conns.items()): + if c.channel is None: + continue + sent = c.taken - c.channel.bufferedAmount \ + if c.channel.readyState == "open" else c.taken + if sent != c.reported: + c.reported = sent + self._put(FLOW, conn, struct.pack(">Q", max(0, sent))) + if tick % STATS_EVERY == 0: + line = sctp_stats.describe(c.pc) + if line != c.stats: + c.stats = line + self._put(STATS, conn, line.encode()) + + +async def _main() -> None: + hello = json.loads(sys.stdin.readline()) + reader, writer = await asyncio.open_connection("127.0.0.1", hello["port"]) + writer.write(bytes.fromhex(hello["token"])) + + root = logging.getLogger() + root.handlers[:] = [_Forward(writer.write)] + root.setLevel(hello.get("level", logging.INFO)) + + # Patches aiortc on import, as the node does for the connections it holds + # itself. + from meshbay_node.transport.webrtc import sctp_fix, sctp_stats # noqa: F401 + + worker = Worker(reader, writer) + worker.configure(hello) + await worker.run() + writer.close() + + +def main() -> None: + try: + asyncio.run(_main()) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/packages/meshbay-node/src/meshbay_node/transport/webrtc/workers.py b/packages/meshbay-node/src/meshbay_node/transport/webrtc/workers.py new file mode 100644 index 0000000..9a00bda --- /dev/null +++ b/packages/meshbay-node/src/meshbay_node/transport/webrtc/workers.py @@ -0,0 +1,392 @@ +"""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 |