path: root/packages/meshbay-node/tests/transfer_probe.py

#!/usr/bin/env python3
"""
Ask a real node for more transfer slots than it has, and watch what it does.

Everything about transfer slots has so far been proved against a `TransferSlots`
object and against sessions with a list where the DataChannel should be. Both
are worth having and neither has ever met a node. This speaks the same MNP over
the same WebRTC DataChannel as a browser, so what it measures is what a member
would get.

Three questions, and only the first is about the cap:

  1. **Is the cap real?** Open more transfers than it allows and count how many
     come back granted. A node that grants everything is a node where none of
     this does anything — which, until MNP 3.0 makes leases compulsory, is also
     what an *old* client gets, so the number here is the difference between
     "built" and "working".
  2. **Does the queue drain?** Close a granted transfer and see whether the
     grant reaches whoever was waiting. The node can be perfectly right about
     who deserves the slot and still never say so; the push is a separate thing
     from the decision, and this is the only place both run.
  3. **Does a slot come back when a peer vanishes?** Drop the connection
     without closing anything — a closed tab, a dead network — and ask a second
     account whether the slot freed. That reclaim is a hook, not a timeout, so
     it should be immediate.

    .venv/bin/python packages/meshbay-node/tests/transfer_probe.py --group <id>
        --want 6         # ask for six slots at once
        --keep-open      # hold them, then look at `meshbay-node transfers`
        --pull 3         # download three files to completion, under a lease
        --pull 3 --parallel   # …at the same time on one connection

**Not collected by pytest** — the filename does not match `test_*.py`, which is
deliberate: this talks to a real hub with real credentials and takes minutes.
It lives here rather than in `QE/` so it survives, because `QE/` is not
versioned and this probe found several defects that no test in the suite could
reach: a cap that was never enforced, a queue that granted a slot and never said
so, and leases that outlived the session holding them.

What it still needs from `QE/`, which stays out of the repo:

  - `QE/deploy/e2e.py` — the second implementation of the client, whose `Client`
    speaks MNP over a real WebRTC DataChannel. Located at run time; the probe
    says so plainly if it is missing rather than failing on an import.
  - `QE/deploy/demo.env` — credentials. Never in the repo, by the same rule.

Nothing here writes to the node: a lease is in-memory state that dies with the
connection, so the worst a failed run leaves behind is a slot the node reclaims
on its own.
"""

import argparse
import asyncio
import sys
import uuid
from pathlib import Path

# `e2e.py` is the MNP client this probe drives, and it lives in QE/, which is
# not versioned (credentials and test artefacts go there by convention). Found
# by walking up to the repo root rather than assumed to be a sibling, and its
# absence is explained rather than raised as an ImportError from four frames
# down.
_QE = Path(__file__).resolve().parents[3] / "QE" / "deploy"
if not (_QE / "e2e.py").exists():
    raise SystemExit(
        f"{__file__.split('/')[-1]} needs QE/deploy/e2e.py, which is not in\n"
        f"this checkout ({_QE} does not have it). QE/ is deliberately not\n"
        f"versioned: it holds credentials and test artefacts. Copy it there, or\n"
        f"run this from a machine that has one.")
sys.path.insert(0, str(_QE))

import httpx   # noqa: E402

from e2e import Client, env   # noqa: E402


def _line(ok: bool, text: str) -> None:
    print(f"  [{'PASS' if ok else 'FAIL'}] {text}")


async def _open_transfer(client: Client, kind: str = "download",
                         nbytes: int = 1 << 30) -> tuple[str, dict]:
    """One `transfer_open`, and the state that comes back."""
    tr = uuid.uuid4().hex
    client.send({"type": "transfer_open", "v": "0.1", "tr": tr,
                 "kind": kind, "bytes": nbytes, "chunks": 1024})
    reply = await client.recv_type("transfer_state", timeout=15)
    return tr, reply


async def pull(client, ack, count: int, parallel: bool = False) -> int:
    """Download files to completion over MNP, under a lease, and say where they
    stop.

    The browser is the hard place to look: a download that freezes near the end
    there could be the service worker's backpressure, the DataChannel, the
    node's own buffer wait, or the lease being revoked underneath it — and the
    interface says the same thing for all four. This client holds no
    SourceBuffer, no service worker and no iframe, so if a file arrives whole
    the node and the transport are cleared and the browser is implicated.
    """
    import time as _t

    # Asked for, not waited for: the node pushes an index when it changes, but
    # a client that has just connected has to request one. Waiting for a push
    # that may never come is a twenty-second timeout that says nothing.
    client.send({"type": "index_sync", "v": "0.1"})
    index = await client.recv_type("index_sync", timeout=30)
    entries = client.open_index(index).get("entries", [])
    big = sorted([e for e in entries if e.get("size", 0) > 50 * 1024 * 1024],
                 key=lambda e: -e["size"])[:count]
    if not big:
        print("no file over 50 MB in this group to pull")
        return 1

    CHUNK = 1024 * 1024
    failures = 0

    if parallel:
        return await pull_together(client, big, CHUNK)

    for entry in big:
        total_chunks = -(-entry["size"] // CHUNK)
        tr, state = await _open_transfer(client, nbytes=entry["size"],
                                         kind="download")
        while state.get("state") == "queued":
            state = await client.recv_type("transfer_state", timeout=120)
        print(f"\n{entry['name'][:52]:<52} {entry['size'] / 1048576:8.1f} MB")

        got = 0
        started = _t.monotonic()
        last_report = started
        try:
            for i in range(total_chunks):
                client.send({"type": "file_req", "v": "0.1",
                             "file_id": entry["id"], "chunk_index": i,
                             "tr": tr})
                msg = await client.recv_type("file_chunk", timeout=90)
                if msg.get("type") == "error":
                    raise RuntimeError(msg.get("detail", "refused"))
                got += len(msg.get("ct") or b"")
                if _t.monotonic() - last_report > 5:
                    last_report = _t.monotonic()
                    print(f"  {got / 1048576:8.1f} MB  chunk {i + 1}/{total_chunks}")
        except Exception as exc:
            pct = 100 * got / max(1, entry["size"])
            print(f"  STOPPED at {got / 1048576:.1f} MB ({pct:.1f}%), "
                  f"chunk of {total_chunks}: {type(exc).__name__}: {exc}")
            failures += 1
            client.send({"type": "transfer_close", "v": "0.1", "tr": tr,
                         "reason": "failed"})
            continue

        secs = _t.monotonic() - started
        ok = got >= entry["size"]
        print(f"  {'COMPLETE' if ok else 'SHORT'} — {got / 1048576:.1f} MB in "
              f"{secs:.0f}s ({got / 1048576 / max(secs, 1):.1f} MB/s)")
        failures += not ok
        client.send({"type": "transfer_close", "v": "0.1", "tr": tr,
                     "reason": "done"})

    print()
    print("every file arrived whole" if not failures
          else f"{failures} file(s) did not arrive whole")
    return 1 if failures else 0


async def pull_together(client, entries, CHUNK) -> int:
    """Every file at once, on one connection, interleaved.

    This is the shape a browser makes and the one a sequential pull cannot
    reproduce: several downloads share a single DataChannel, each keeping a
    window of chunk requests in flight, so the node's send buffer is under
    pressure from all of them at once and every reply waits behind the others.
    A download that arrives whole on its own can still stall here.

    Replies are matched by (file_id, chunk_index) rather than by arrival order,
    because with several downloads in flight arrival order means nothing —
    which is the same reason `req_id` exists.
    """
    import time as _t

    WINDOW = 8              # PIPELINE_WINDOW in file-utils.js
    state = {}
    for e in entries:
        tr, st = await _open_transfer(client, nbytes=e["size"], kind="download")
        while st.get("state") == "queued":
            st = await client.recv_type("transfer_state", timeout=180)
        state[e["id"]] = {"entry": e, "tr": tr, "sent": 0, "got": 0,
                          "bytes": 0, "total": -(-e["size"] // CHUNK)}
        print(f"{e['name'][:52]:<52} {e['size'] / 1048576:8.1f} MB")

    def fire():
        for st in state.values():
            while st["sent"] < st["total"] and st["sent"] - st["got"] < WINDOW:
                client.send({"type": "file_req", "v": "0.1",
                             "file_id": st["entry"]["id"],
                             "chunk_index": st["sent"], "tr": st["tr"]})
                st["sent"] += 1

    fire()
    started = last = _t.monotonic()
    stalled = None
    while any(st["got"] < st["total"] for st in state.values()):
        try:
            msg = await client.recv_type("file_chunk", timeout=45)
        except Exception as exc:
            stalled = f"{type(exc).__name__}: {exc}"
            break
        if msg.get("type") == "error":
            stalled = f"node refused: {msg.get('detail')}"
            break
        st = state.get(msg.get("file_id"))
        if st is None:
            continue
        st["got"] += 1
        st["bytes"] += len(msg.get("ct") or b"")
        fire()
        if _t.monotonic() - last > 5:
            last = _t.monotonic()
            print("  " + " | ".join(
                f"{s['entry']['name'][:14]:<14} {s['got']:>4}/{s['total']}"
                for s in state.values()))

    print()
    failures = 0
    for st in state.values():
        done = st["got"] >= st["total"]
        failures += not done
        print(f"  {'COMPLETE' if done else 'STOPPED '}  "
              f"{st['entry']['name'][:44]:<44} "
              f"{st['bytes'] / 1048576:8.1f} MB  "
              f"chunk {st['got']}/{st['total']}")
        client.send({"type": "transfer_close", "v": "0.1", "tr": st["tr"],
                     "reason": "done" if done else "failed"})
    if stalled:
        print(f"\n  stalled after {_t.monotonic() - started:.0f}s: {stalled}")
    print()
    print("every file arrived whole" if not failures
          else f"{failures} of {len(state)} did not arrive whole")
    return 1 if failures else 0


async def probe(args) -> int:
    cfg = env()
    hub = cfg["HUB_URL"]
    failures = 0

    async with httpx.AsyncClient(timeout=30) as http:
        alice = Client(hub, cfg["NODE_USER"], cfg["NODE_PASS"])
        await alice.login(http)

        # Which node is serving this group, resolved the way stream_probe.py
        # does it. `connect()` needs the node id as well as the group: the hub
        # relays signalling to one node, and a group may be hosted by more than
        # one.
        groups = (await http.get(f"{hub}/v1/groups/mine",
                                 headers=alice.auth)).json()["groups"]
        wanted = [g for g in groups
                  if g["id"] == args.group
                  or g["id"].startswith(args.group)
                  or args.group.lower() in g["name"].lower()]
        if not wanted:
            print(f"no group of yours matches {args.group!r}")
            return 1
        group = wanted[0]
        nodes = (await http.get(f"{hub}/v1/groups/{group['id']}/nodes",
                                headers=alice.auth)).json()["nodes"]
        if not nodes:
            print(f"no node online for {group['name']} — start it and retry")
            return 1
        node_id = nodes[0]["node_id"]
        print(f"group  : {group['name']} ({group['id'][:8]})")
        print(f"node   : {node_id[:12]}\n")

        ack = await alice.connect(http, group["id"], node_id)

        limits = ack.get("transfer_limits")
        if limits is None:
            print("This node does not hand out transfer slots — it predates "
                  "them, or the handshake ack lost the field. Nothing below "
                  "can be measured.")
            await alice.close()
            return 1
        cap = int(limits.get("download") or 0)
        print(f"node reports this member may run {cap} download(s) at once\n")

        if args.pull:
            return await pull(alice, ack, args.pull, args.parallel)

        # ── 1. is the cap real ────────────────────────────────────────────
        #
        # The baseline first. A node that is already serving somebody grants
        # this probe fewer slots than its cap, entirely correctly — and an
        # earlier version reported that as two failures, which is a probe
        # lying about a node that was right. Seen for real: a previous run of
        # this script had crashed before closing, and its leases were still
        # held. So the first reply is read for what the node says is already
        # in use, and the run stops rather than measuring against a moving
        # floor.
        want = args.want or (cap + 2)
        opened = [await _open_transfer(alice)]
        first = opened[0][1]
        if first.get("state") != "granted" or first.get("used", 1) != 1:
            print(f"this node is not idle: it reports {first.get('used')} of "
                  f"{first.get('cap')} slots already used by this member, and "
                  f"{first.get('node_used')} of {first.get('node_cap')} "
                  f"node-wide.\nWait for it to settle (or `meshbay-node "
                  f"transfers show` to see what is holding them) and run again "
                  f"— the cap cannot be measured against a moving floor.")
            await alice.close()
            return 1
        for _ in range(want - 1):
            opened.append(await _open_transfer(alice))
        granted = [r for _, r in opened if r.get("state") == "granted"]
        queued = [r for _, r in opened if r.get("state") == "queued"]
        print(f"asked for {want}: {len(granted)} granted, {len(queued)} queued")
        ok = len(granted) == cap and len(queued) == want - cap
        _line(ok, f"the cap is enforced ({len(granted)} granted against a cap "
                  f"of {cap})")
        failures += not ok

        if queued:
            positions = [r.get("ahead") for r in queued]
            ok = positions == sorted(positions) and positions[0] == 0
            _line(ok, f"queue positions are handed out in order: {positions}")
            failures += not ok

        if args.keep_open:
            print("\nholding them. Look at the node with:\n"
                  "    meshbay-node transfers show\n"
                  "Ctrl-C when done — every slot is released by the "
                  "disconnection alone.")
            try:
                await asyncio.Event().wait()
            except (KeyboardInterrupt, asyncio.CancelledError):
                pass
            await alice.close()
            return 0

        # ── 2. does the queue drain ───────────────────────────────────────
        if queued:
            first_tr = opened[0][0]
            alice.send({"type": "transfer_close", "v": "0.1", "tr": first_tr,
                        "reason": "done"})
            # Two messages come back: the close, and the grant it produced.
            # Which order is not promised, so both are collected.
            seen = []
            for _ in range(2):
                try:
                    seen.append(await alice.recv_type("transfer_state",
                                                      timeout=15))
                except asyncio.TimeoutError:
                    break
            promoted = [m for m in seen if m.get("state") == "granted"]
            ok = bool(promoted)
            _line(ok, "closing a transfer granted the slot to the next in queue")
            failures += not ok
            if not ok:
                print("        the node decided correctly and never said so — "
                      "the client would sit at 'waiting' for ever")

        # ── 3. does a vanished peer give its slots back ───────────────────
        #
        # No second account needed, and that is not a compromise: a lease
        # belongs to the *connection*, so a reconnecting client is a new session
        # to the node. If the old one's leases were not released they still
        # count against this member's cap, and the reconnect finds nothing free
        # — which makes this the same check, on any group, without depending on
        # who else happens to be a member.
        #
        # Dropped without closing anything: a shut tab, a dead network. The
        # reclaim is a hook on the connection, not a timeout, so it should be
        # immediate rather than two minutes away.
        await alice.close()
        await asyncio.sleep(1.5)

        again = Client(hub, cfg["NODE_USER"], cfg["NODE_PASS"])
        await again.login(http)
        await again.connect(http, group["id"], node_id)
        _, reborn = await _open_transfer(again)
        ok = reborn.get("state") == "granted" and reborn.get("used") == 1
        _line(ok, "the slots came back when the peer vanished without closing")
        failures += not ok
        if not ok:
            print(f"        the node still counts {reborn.get('used')} of "
                  f"{reborn.get('cap')} against this member — the old session's "
                  f"leases outlived it, and only the idle sweep will free them")
        await again.close()

    print()
    if failures:
        print(f"{failures} check(s) failed — do not make leases compulsory yet")
    else:
        print("all checks passed")
    return 1 if failures else 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--group", required=True, help="group id to connect to")
    ap.add_argument("--want", type=int, default=0,
                    help="how many transfers to open at once (default: cap + 2)")
    ap.add_argument("--keep-open", action="store_true",
                    help="hold the transfers so the node can be inspected")
    ap.add_argument("--pull", type=int, default=0, metavar="N",
                    help="actually download N files to completion, under a "
                         "lease, and report where they stop")
    ap.add_argument("--parallel", action="store_true",
                    help="pull them at the same time on one connection, the "
                         "way a browser does — which is when it goes wrong")
    return asyncio.run(probe(ap.parse_args()))


if __name__ == "__main__":
    raise SystemExit(main())