d8208685a2
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>
364 lines
15 KiB
Python
364 lines
15 KiB
Python
"""
|
|
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())
|