Initial commit: multi-peer torrent download engine

Reactor/loop-pool engine with TCP/µTP/MSE transports, per-connection
pipelining, priority-driven piece selection with endgame, and the Python
FFI test harness.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

harness/seed_server.py

@@ -0,0 +1,364 @@
+"""
+seed_server.py - A lightweight emulated BitTorrent client (seed side) for
+load-testing the peer/engine against "as many peers as possible".
+
+Rather than spinning up N heavyweight libtorrent sessions, this serves one
+torrent's data from a single asyncio event loop across many listener sockets.
+Each listener is a distinct endpoint, so from the engine's point of view each is
+a separate peer: one `add_peer(tid, ip, port)` per listener.
+
+It speaks the plaintext BitTorrent v1 wire protocol a real client would on the
+seed side: validates the handshake + info-hash, advertises a full bitfield,
+unchokes, and answers `request` messages with `piece` data read (mmap'd) from
+disk. That is exactly the subset the engine drives, and it is plaintext because
+the engine does not negotiate MSE encryption yet.
+
+Usage
+-----
+Serve an existing torrent (data already on disk under --data):
+    python harness/seed_server.py some.torrent --data /path/to/datadir --peers 64
+
+Generate a random test torrent, serve it, and print the .torrent path:
+    python harness/seed_server.py --generate 256M --peers 64 --out /tmp/seedtest
+
+Generate + serve + drive the engine against every peer and report throughput:
+    python harness/seed_server.py --generate 256M --peers 64 --out /tmp/seedtest --self-test
+
+In a test, use the SeedSwarm class directly to start listeners in-process and
+read `swarm.endpoints`.
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import mmap
+import os
+import secrets
+import struct
+import threading
+import time
+
+PSTR = b"BitTorrent protocol"
+HANDSHAKE_LEN = 68
+
+# Wire message ids.
+MSG_CHOKE, MSG_UNCHOKE, MSG_INTERESTED, MSG_NOT_INTERESTED = 0, 1, 2, 3
+MSG_HAVE, MSG_BITFIELD, MSG_REQUEST, MSG_PIECE, MSG_CANCEL = 4, 5, 6, 7, 8
+
+DRAIN_HIGH_WATER = 1 << 20  # let blocks queue up, apply backpressure past 1 MiB
+MAX_BLOCK = 1 << 17         # reject absurd request lengths (128 KiB)
+
+
+# --------------------------------------------------------------------------- #
+# Data source: read (offset, length) from the torrent's concatenated files.
+# --------------------------------------------------------------------------- #
+class DataSource:
+    """Maps the linear piece space onto one or more on-disk files (BT v1 lays
+    files out back-to-back). Files are mmap'd for cheap repeated reads."""
+
+    def __init__(self, files: list[tuple[str, int]]):
+        self._maps = []  # (global_offset, size, mmap_or_none)
+        self._handles = []
+        off = 0
+        for path, size in files:
+            mm = None
+            if size > 0:
+                fh = open(path, "rb")
+                self._handles.append(fh)
+                mm = mmap.mmap(fh.fileno(), size, prot=mmap.PROT_READ)
+            self._maps.append((off, size, mm))
+            off += size
+        self.total = off
+        self._single = self._maps[0][2] if len(self._maps) == 1 else None
+
+    def read(self, offset: int, length: int) -> bytes:
+        if self._single is not None:                 # fast path: one file
+            return self._single[offset:offset + length]
+        out = bytearray()
+        remaining = length
+        for foff, size, mm in self._maps:
+            if remaining <= 0:
+                break
+            if offset >= foff + size or offset < foff:
+                continue
+            local = offset - foff
+            take = min(size - local, remaining)
+            out += mm[local:local + take]
+            offset += take
+            remaining -= take
+        return bytes(out)
+
+    def close(self):
+        for _, _, mm in self._maps:
+            if mm is not None:
+                mm.close()
+        for fh in self._handles:
+            fh.close()
+
+
+def _full_bitfield(num_pieces: int) -> bytes:
+    nbytes = (num_pieces + 7) // 8
+    bf = bytearray(b"\xff" * nbytes)
+    rem = num_pieces & 7
+    if rem:                                          # clear pad bits past the end
+        bf[-1] = (0xFF << (8 - rem)) & 0xFF
+    return bytes(bf)
+
+
+def _msg(mid: int, payload: bytes = b"") -> bytes:
+    return struct.pack(">IB", 1 + len(payload), mid) + payload
+
+
+# --------------------------------------------------------------------------- #
+# The swarm: many listeners sharing one event loop, on a background thread.
+# --------------------------------------------------------------------------- #
+class SeedSwarm:
+    def __init__(self, info_hash: bytes, piece_length: int, num_pieces: int,
+                 source: DataSource, host: str = "127.0.0.1"):
+        self.info_hash = info_hash
+        self.piece_length = piece_length
+        self.num_pieces = num_pieces
+        self.source = source
+        self.host = host
+        self.bitfield = _full_bitfield(num_pieces)
+        self.endpoints: list[tuple[str, int]] = []
+        self.served_bytes = 0
+        self._loop = asyncio.new_event_loop()
+        self._servers: list[asyncio.AbstractServer] = []
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    # -- connection handler (one per accepted peer) ------------------------- #
+    async def _handle(self, reader: asyncio.StreamReader,
+                      writer: asyncio.StreamWriter):
+        try:
+            hs = await reader.readexactly(HANDSHAKE_LEN)
+            if hs[1:20] != PSTR or hs[28:48] != self.info_hash:
+                writer.close()
+                return
+            peer_id = b"-SD0001-" + secrets.token_bytes(12)
+            writer.write(bytes([len(PSTR)]) + PSTR + b"\x00" * 8 +
+                         self.info_hash + peer_id)
+            writer.write(_msg(MSG_BITFIELD, self.bitfield))
+            writer.write(_msg(MSG_UNCHOKE))
+            await writer.drain()
+
+            while True:
+                (length,) = struct.unpack(">I", await reader.readexactly(4))
+                if length == 0:
+                    continue                          # keep-alive
+                body = await reader.readexactly(length)
+                mid = body[0]
+                if mid == MSG_REQUEST and length >= 13:
+                    index, begin, blen = struct.unpack(">III", body[1:13])
+                    if blen > MAX_BLOCK:
+                        continue
+                    data = self.source.read(index * self.piece_length + begin, blen)
+                    writer.write(struct.pack(">IB", 9 + len(data), MSG_PIECE) +
+                                 struct.pack(">II", index, begin) + data)
+                    self.served_bytes += len(data)
+                    if writer.transport.get_write_buffer_size() > DRAIN_HIGH_WATER:
+                        await writer.drain()
+                # interested / not-interested / cancel / choke: nothing to do —
+                # we are already unchoked and answer requests as they arrive.
+        except (asyncio.IncompleteReadError, ConnectionResetError,
+                BrokenPipeError, ConnectionError):
+            pass
+        finally:
+            try:
+                writer.close()
+            except Exception:
+                pass
+
+    async def _make_servers(self, count: int):
+        servers, endpoints = [], []
+        for _ in range(count):
+            srv = await asyncio.start_server(self._handle, self.host, 0)
+            servers.append(srv)
+            endpoints.append((self.host, srv.sockets[0].getsockname()[1]))
+        return servers, endpoints
+
+    def _run(self):
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def start(self, count: int):
+        self._thread.start()
+        fut = asyncio.run_coroutine_threadsafe(self._make_servers(count), self._loop)
+        self._servers, self.endpoints = fut.result(timeout=15)
+        return self.endpoints
+
+    def stop(self):
+        def _shutdown():
+            for srv in self._servers:
+                srv.close()
+            self._loop.stop()
+        if self._loop.is_running():
+            self._loop.call_soon_threadsafe(_shutdown)
+        self._thread.join(timeout=5)
+        try:
+            self._loop.close()
+        except Exception:
+            pass
+        self.source.close()
+
+
+# --------------------------------------------------------------------------- #
+# Torrent loading / generation.
+# --------------------------------------------------------------------------- #
+def load_torrent(torrent_path: str, data_dir: str):
+    """Return (info_hash, piece_length, num_pieces, total_size, DataSource)."""
+    from torrent_meta import load_torrent as parse_torrent
+    tf = parse_torrent(torrent_path)
+    files = []
+    for path, size in tf.files:
+        files.append((os.path.join(data_dir, path), size))
+    missing = [p for p, _ in files if not os.path.exists(p)]
+    if missing:
+        raise FileNotFoundError(f"data files not found under {data_dir}: {missing}")
+    meta = tf.metadata
+    return (meta.info_hash, meta.piece_length, meta.num_pieces,
+            meta.total_size, DataSource(files))
+
+
+def generate_torrent(out_dir: str, size: int, piece_size: int = 256 * 1024):
+    """Create a random data file + .torrent under out_dir. Returns
+    (torrent_path, info_hash, piece_length, num_pieces, total_size, DataSource)."""
+    import libtorrent as lt
+    os.makedirs(out_dir, exist_ok=True)
+    data_path = os.path.join(out_dir, "data.bin")
+    with open(data_path, "wb") as f:
+        rem = size
+        while rem > 0:
+            n = min(rem, 8 << 20)
+            f.write(os.urandom(n))
+            rem -= n
+    fs = lt.file_storage()
+    lt.add_files(fs, data_path)
+    t = lt.create_torrent(fs, piece_size=piece_size)
+    t.set_priv(False)
+    lt.set_piece_hashes(t, out_dir)
+    torrent_path = os.path.join(out_dir, "test.torrent")
+    with open(torrent_path, "wb") as f:
+        f.write(lt.bencode(t.generate()))
+    ih, pl, npc, total, src = load_torrent(torrent_path, out_dir)
+    return torrent_path, ih, pl, npc, total, src
+
+
+def parse_size(s: str) -> int:
+    s = s.strip().upper()
+    mult = 1
+    if s and s[-1] in "KMG":
+        mult = {"K": 1024, "M": 1024**2, "G": 1024**3}[s[-1]]
+        s = s[:-1]
+    return int(float(s) * mult)
+
+
+# --------------------------------------------------------------------------- #
+# Optional: drive the engine against every peer and verify + report rate.
+# --------------------------------------------------------------------------- #
+def _self_test(torrent_path: str, data_dir: str, endpoints, timeout: float):
+    import hashlib
+    from engine_ffi import Engine, EngineConfig, STATE_ERROR, ERROR_NAMES
+    from harness import load_metadata
+
+    meta = load_metadata(torrent_path)
+    lib = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                       "build", "libtorrentpeer.so")
+    loops = min(8, max(1, (os.cpu_count() or 1)))
+    with Engine(EngineConfig(loop_count=loops, slots_per_loop=4096,
+                             max_pipeline=1024), lib_path=lib) as eng:
+        tid = eng.add_torrent(meta.info_hash, b"-PC0001-" + secrets.token_bytes(12),
+                              meta.piece_length, meta.total_size, meta.num_pieces)
+        eng.set_priorities(tid, [1] * meta.num_pieces)
+        for ip, port in endpoints:
+            eng.add_peer(tid, ip, port)
+
+        buffers = [bytearray(meta.piece_len(i)) for i in range(meta.num_pieces)]
+        received = [0] * meta.num_pieces
+        done = bytearray(meta.num_pieces)
+        done_count = 0
+        t0 = time.time()
+        deadline = t0 + timeout
+        while done_count < meta.num_pieces:
+            st = eng.status(tid)
+            if st.state == STATE_ERROR:
+                raise RuntimeError(f"engine error: {ERROR_NAMES[st.error]}")
+            descs = eng.poll_ready()
+            if not descs:
+                eng.wait(200)
+                if time.time() > deadline:
+                    raise TimeoutError(f"stalled at {done_count}/{meta.num_pieces}")
+                continue
+            for x in descs:
+                buf = buffers[x.piece]
+                buf[x.begin:x.begin + x.len] = eng.block_data(x.loop, x.slot, x.len)
+                eng.release(x.loop, x.slot)
+                received[x.piece] += x.len
+                if not done[x.piece] and received[x.piece] >= meta.piece_len(x.piece):
+                    if hashlib.sha1(bytes(buf)).digest() != meta.piece_hashes[x.piece]:
+                        raise ValueError(f"piece {x.piece} hash mismatch")
+                    done[x.piece] = 1
+                    done_count += 1
+                    eng.set_priority(tid, x.piece, 0)
+        dt = time.time() - t0
+        mb = meta.total_size / 1e6
+        print(f"self-test OK: {mb:.1f} MB from {len(endpoints)} peers "
+              f"in {dt:.2f}s = {mb/dt:.1f} MB/s (peers={eng.status(tid).peers})")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("torrent", nargs="?", help="path to an existing .torrent")
+    ap.add_argument("--data", help="directory (or file) holding the torrent's data")
+    ap.add_argument("--generate", metavar="SIZE",
+                    help="generate a random torrent of this size (e.g. 256M, 1G)")
+    ap.add_argument("--out", default="/tmp/seedtest",
+                    help="output dir for --generate (default: /tmp/seedtest)")
+    ap.add_argument("--peers", type=int, default=32, help="number of listeners")
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--piece-size", type=int, default=256 * 1024)
+    ap.add_argument("--ports-file", help="write the endpoints (one ip:port/line) here")
+    ap.add_argument("--self-test", action="store_true",
+                    help="drive the engine against all peers, verify, report rate")
+    ap.add_argument("--timeout", type=float, default=300.0)
+    args = ap.parse_args()
+
+    torrent_path = args.torrent
+    if args.generate:
+        size = parse_size(args.generate)
+        torrent_path, ih, pl, npc, total, src = generate_torrent(
+            args.out, size, args.piece_size)
+        data_dir = args.out
+        print(f"generated {torrent_path}  ({total} bytes, {npc} pieces x {pl})")
+    else:
+        if not torrent_path or not args.data:
+            ap.error("provide a .torrent and --data, or use --generate SIZE")
+        data_dir = args.data if os.path.isdir(args.data) else os.path.dirname(args.data)
+        ih, pl, npc, total, src = load_torrent(torrent_path, data_dir)
+
+    swarm = SeedSwarm(ih, pl, npc, src, host=args.host)
+    swarm.start(args.peers)
+    print(f"seeding {total} bytes on {len(swarm.endpoints)} peers "
+          f"({args.host}:{swarm.endpoints[0][1]}..{swarm.endpoints[-1][1]})")
+    if args.ports_file:
+        with open(args.ports_file, "w") as f:
+            f.writelines(f"{ip}:{port}\n" for ip, port in swarm.endpoints)
+        print(f"endpoints written to {args.ports_file}")
+
+    try:
+        if args.self_test:
+            _self_test(torrent_path, data_dir, swarm.endpoints, args.timeout)
+        else:
+            print("serving; press Ctrl-C to stop")
+            while True:
+                time.sleep(1.0)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        swarm.stop()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())