diff options
Diffstat (limited to 'packages/meshbay-node/src/meshbay_node/transport')
| -rw-r--r-- | packages/meshbay-node/src/meshbay_node/transport/webrtc_server.py | 249 |
1 files changed, 248 insertions, 1 deletions
diff --git a/packages/meshbay-node/src/meshbay_node/transport/webrtc_server.py b/packages/meshbay-node/src/meshbay_node/transport/webrtc_server.py index 33f5474..164cc62 100644 --- a/packages/meshbay-node/src/meshbay_node/transport/webrtc_server.py +++ b/packages/meshbay-node/src/meshbay_node/transport/webrtc_server.py @@ -127,6 +127,8 @@ from meshbay_node.transport.wire import index_sync_message from meshbay_node.indexer import GroupIndex from meshbay_node.indexer.indexer import DirectoryIndexer from meshbay_node import linkpreview, ops, platform +from meshbay_node import transfers as transfers_mod +from meshbay_node.transfers import TransferSlots # Re-imported under its original name: every call site and existing test in # this module still refers to it as `_probe_video`. The implementation lives # in media_probe.py so the indexer package (imported just above) can call it @@ -258,6 +260,11 @@ STREAM_CREDIT_TIMEOUT = 120 # How often that budget is re-examined. A viewer who left stops being # charged for a slot within this, rather than within the timeout. STREAM_CREDIT_POLL = 3 +# How often transfer leases are swept. Nothing depends on it being +# prompt -- the session teardown is the reclaim that matters and is +# immediate; this catches peers that vanished without the connection +# noticing, so it trades latency for a timer that hardly ever runs. +TRANSFER_SWEEP_SECS = 15 def _pack(obj: dict) -> bytes: @@ -524,6 +531,10 @@ class WebRTCPeerSession: self._spawn(self._do_link_preview_request(msg)) elif mtype == MNP.PING: self._do_ping(msg) + elif mtype == MNP.TRANSFER_OPEN: + self._do_transfer_open(msg) + elif mtype == MNP.TRANSFER_CLOSE: + self._do_transfer_close(msg) elif mtype == MNP.FILE_UPLOAD: self._do_file_upload(msg) elif mtype == MNP.DIR_CREATE: @@ -3309,6 +3320,187 @@ class WebRTCPeerSession: "total_bytes": progress.total_bytes, } + # ── Transfer slots ─────────────────────────────────────────────────────── + + def _slots(self) -> "TransferSlots": + """The node's transfer pools, shared across every peer and every group. + + On the transport context, not the session: it counts the node's + transfers, not one browser's. Built once, for the same reason the + transcode semaphore is — rebuilding it per call would hand every caller + its own budget and cap nothing at all. + """ + slots = self._ctx.get("_transfer_slots") + if slots is None: + slots = TransferSlots() + n = self._ctx.get("max_concurrent_downloads") + u = self._ctx.get("max_concurrent_uploads") + if n: + slots.caps[transfers_mod.DOWNLOAD] = int(n) + if u: + slots.caps[transfers_mod.UPLOAD] = int(u) + self._ctx["_transfer_slots"] = slots + log.info("transfer: %s", slots.summary()) + return slots + + def _transfer_state_msg(self, lease, state: str, reason: str = "") -> dict: + slots = self._slots() + out = { + "type": MNP.TRANSFER_STATE, + "v": MNP_VERSION, + "tr": lease.tr, + "state": state, + "kind": lease.kind, + "used": slots.member_in_use(lease.kind, lease.member), + "cap": slots.per_member.get(lease.kind, + transfers_mod.DEFAULT_MAX_PER_MEMBER), + "node_used": slots.in_use(lease.kind), + "node_cap": slots.caps.get(lease.kind, + transfers_mod.DEFAULT_MAX_CONCURRENT), + } + if state == "queued": + out["ahead"] = slots.ahead_of(lease) + if reason: + out["reason"] = reason + return out + + def _notify_transfer(self, lease, state: str, reason: str = "") -> None: + """Push a lease's state to the connection that owns it. + + By session key, never by account: a lease belongs to one connection, and + telling a member's other device that *its* transfer was granted is how a + queue starts lying. + """ + session = self._peer_registry().get(lease.session_key) + for candidate in ([session] if session else + self._sessions_everywhere(lease.session_key)): + try: + candidate._send(self._transfer_state_msg(lease, state, reason)) + except Exception: + pass + + def _sessions_everywhere(self, session_key: str) -> list["WebRTCPeerSession"]: + """The session with this key, whichever group it is in. + + `_peer_registry` is per group (finding H1) and the pools are node-wide, + so a slot freed in one group can grant one in another: the peer to tell + is not necessarily in this session's own registry. + """ + groups = self._ctx.get("groups") + registries = ([g.get("_peers", {}) for g in groups.values()] + if groups else [self._ctx.get("_peers", {})]) + return [reg[session_key] for reg in registries if session_key in reg] + + def _announce(self, granted: list, ended: list | None = None) -> None: + for lease, reason in (ended or []): + self._notify_transfer( + lease, "queued" if lease.state == "queued" else "closed", reason) + for lease in granted: + self._notify_transfer(lease, "granted") + + def _do_transfer_open(self, msg: dict) -> None: + tr = str(msg.get("tr") or "")[:64] + kind = str(msg.get("kind") or transfers_mod.DOWNLOAD) + if not tr: + self._send({"type": "error", "detail": "Missing transfer id", + "code": "bad_transfer_id"}) + return + slots = self._slots() + try: + nbytes = int(msg.get("bytes") or 0) + chunks = int(msg.get("chunks") or 0) + except (TypeError, ValueError): + self._send({"type": "error", "detail": "Invalid transfer size", + "code": "bad_transfer_size", "tr": tr}) + return + lease, err = slots.open( + tr=tr, kind=kind, session_key=self._registry_key, + user_id=self._user_id or "", group_id=self._group_id or "", + bytes=nbytes, chunks=chunks) + if err: + self._send({"type": "error", "detail": err, "code": err, "tr": tr}) + return + self._send(self._transfer_state_msg(lease, lease.state)) + log.debug("transfer: open %s %s -> %s (%s)", + kind, tr[:8], lease.state, slots.summary()) + self._ensure_transfer_sweeper() + + def _do_transfer_close(self, msg: dict) -> None: + tr = str(msg.get("tr") or "")[:64] + reason = str(msg.get("reason") or transfers_mod.REASON_DONE)[:32] + slots = self._slots() + held = slots.leases.get(tr) + if held is not None and held.session_key != self._registry_key: + # Closing somebody else's transfer would be a denial of service one + # random id away. + self._send({"type": "error", "detail": "not_your_transfer", + "code": "not_your_transfer", "tr": tr}) + return + lease, granted = slots.close(tr, reason) + if lease is not None: + self._send(self._transfer_state_msg(lease, "closed", reason)) + self._announce(granted) + + def _release_transfers(self) -> None: + """Give back everything this connection held. Called from teardown.""" + slots = self._ctx.get("_transfer_slots") + if slots is None: + return + gone, granted = slots.release_session(self._registry_key) + if gone: + log.info("transfer: session gone, released %d (%s)", + len(gone), slots.summary()) + self._announce(granted) + + def _ensure_transfer_sweeper(self) -> None: + """Start the maintenance task, once, and only while it has work. + + It reclaims what a session teardown cannot see — a grant nobody took up, + a transfer that went quiet — and logs the one line that answers "was + this peer ever in a queue" when somebody reports a stuck transfer. It + stops when the last lease goes, so an idle node runs no timer. + """ + running = self._ctx.get("_transfer_sweeper") + if running is not None and not running.done(): + return + + ctx = self._ctx + + async def _sweep_loop() -> None: + while True: + await asyncio.sleep(TRANSFER_SWEEP_SECS) + slots = ctx.get("_transfer_slots") + if slots is None or not slots.leases: + return + ended, granted = slots.sweep() + for lease, reason in ended: + log.info("transfer: reclaimed %s (%s)", lease.tr[:8], reason) + self._announce(granted, ended) + log.debug("transfer: %s", slots.summary()) + + # Deliberately NOT `self._spawn`, which is otherwise the only way to + # start a task here. `_spawn` ties a task to *this session's* set, and + # `shutdown_tasks` cancels those when the peer leaves — so the sweeper + # would die with whichever connection happened to open the first + # transfer, and every other peer's abandoned lease would then never be + # reclaimed. It belongs to the node, so the strong reference that keeps + # it off the garbage collector lives on the transport context; the rule + # `_spawn` exists for (asyncio holds only a weak reference) is satisfied + # by that reference, not by which set it is in. + task = asyncio.ensure_future(_sweep_loop()) + ctx["_transfer_sweeper"] = task + + def _finished(done: asyncio.Task) -> None: + if ctx.get("_transfer_sweeper") is done: + ctx["_transfer_sweeper"] = None + if not done.cancelled() and done.exception() is not None: + # Nothing awaits this task, so an exception here would otherwise + # be swallowed and idle leases would silently stop being + # reclaimed — the failure mode is a node that fills up over days. + log.error("transfer: sweeper died: %r", done.exception()) + + task.add_done_callback(_finished) + def _register_peer(self) -> None: """Add this connection to its group's peer set. @@ -5507,6 +5699,12 @@ class WebRTCPeerSession: here: cancelling the task runs the exit of its `async with sem`. """ self._stop_stream() + # Before the tasks are cancelled: a lease is not held by a task, so + # nothing else would give it back, and this hook is the one place every + # way of walking away arrives at (see the connectionstatechange handler, + # which calls it for a closed tab, a quit browser and a dead network + # alike). + self._release_transfers() for task in list(self._tasks): task.cancel() if self._tasks: @@ -5514,6 +5712,7 @@ class WebRTCPeerSession: async def close(self) -> None: self._audit("disconnect") + self._release_transfers() if self._user_id: self._unregister_peer() await self.shutdown_tasks() @@ -5620,7 +5819,9 @@ class WebRTCTransport: self._stun = stun_servers or list(DEFAULT_STUN_SERVERS) self._sessions: dict[str, WebRTCPeerSession] = {} - def set_capacity(self, *, max_concurrent_streams: int | None = None) -> dict: + def set_capacity(self, *, max_concurrent_streams: int | None = None, + max_concurrent_downloads: int | None = None, + max_concurrent_uploads: int | None = None) -> dict: """Resize a live pool without restarting the daemon. `ops.set_node_settings` used to do this by assigning @@ -5660,8 +5861,54 @@ class WebRTCTransport: log.info("stream: capacity %s -> %d (no pool built yet)", before, n) changed["max_concurrent_streams"] = n + + pools = {} + if max_concurrent_downloads is not None: + pools[transfers_mod.DOWNLOAD] = int(max_concurrent_downloads) + if max_concurrent_uploads is not None: + pools[transfers_mod.UPLOAD] = int(max_concurrent_uploads) + for key, value in pools.items(): + if value < 1: + raise ValueError(f"max_concurrent_{key}s must be positive") + if pools: + # Kept on the context whether or not a pool exists yet: the pools + # are built on the first transfer, and would otherwise come up with + # the defaults after an operator had already changed them. + for key, value in pools.items(): + self._ctx[f"max_concurrent_{key}s"] = value + changed[f"max_concurrent_{key}s"] = value + slots = self._ctx.get("_transfer_slots") + if slots is not None: + granted = slots.set_caps(node=pools) + log.info("transfer: capacity now %s (%d started at once)", + slots.summary(), len(granted)) + # Raising a cap can start queued transfers immediately, and the + # peers waiting on them have to be told: a grant nobody hears + # about is the "stuck at waiting" report this design exists to + # prevent. + for lease in granted: + self._notify_granted(lease) return changed + def _notify_granted(self, lease) -> None: + """Tell the connection that owns `lease` it may start. + + On the transport rather than the session because a cap change has no + session behind it — it arrives from the loopback API. + """ + groups = self._ctx.get("groups") + registries = ([g.get("_peers", {}) for g in groups.values()] + if groups else [self._ctx.get("_peers", {})]) + for reg in registries: + session = reg.get(lease.session_key) + if session is not None: + try: + session._send( + session._transfer_state_msg(lease, "granted")) + except Exception: + pass + return + async def handle_offer( self, offer_sdp: str, peer_id: str, ) -> tuple[str, list[dict]]: |