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>
This commit is contained in:
ookami125
commit
d8208685a2
55 changed files with 9989 additions and 0 deletions
220
tests/test_endgame.py
Normal file
220
tests/test_endgame.py
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
"""
|
||||
Regression test for engine-side endgame (scheduler.c).
|
||||
|
||||
The bug: a piece is claimed loop-wide (tor->requested[i]=1) by one connection,
|
||||
so no other peer will request it. If that peer stays alive but stops delivering
|
||||
(slow/snubbing seed), the piece never completes and healthy seeds sit idle even
|
||||
though they have the data -- the classic "a few pieces never finish" tail stall.
|
||||
|
||||
This test reproduces it deterministically with a raw "stalling" seed that serves
|
||||
every block EXCEPT the last block of each piece (so it keeps its claims but never
|
||||
finishes them, and the connection stays alive so the claim is never released). A
|
||||
second, healthy seed is added afterwards. Only engine endgame -- letting an idle
|
||||
peer race blocks of an already-claimed piece -- can finish the download, so a
|
||||
regression here turns into a TimeoutError from drive_engine().
|
||||
|
||||
Run with: python tests/test_endgame.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, os.path.join(ROOT, "harness"))
|
||||
|
||||
from harness import load_metadata # noqa: E402
|
||||
from engine_ffi import Engine, EngineConfig # noqa: E402
|
||||
from test_localseed import ensure_built, make_torrent # noqa: E402
|
||||
from test_engine import TorrentDrive, drive_engine, make_peer_id # noqa: E402
|
||||
from test_mockpeer import _recv_exact, _msg, _piece_msg # noqa: E402
|
||||
|
||||
LIB = os.path.join(ROOT, "build", "libtorrentpeer.so")
|
||||
BLOCK = 16384
|
||||
|
||||
|
||||
class StallingSeed(threading.Thread):
|
||||
"""Seed that has every piece but never serves the LAST block of any piece.
|
||||
|
||||
It keeps reading (and ignoring) the withheld requests, so the connection
|
||||
stays healthy and its piece claims are never released -- exactly the state
|
||||
that wedges the tail without endgame.
|
||||
"""
|
||||
|
||||
def __init__(self, data, meta):
|
||||
super().__init__(daemon=True)
|
||||
self.data = data
|
||||
self.meta = meta
|
||||
self.served = 0
|
||||
self.srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
self.srv.bind(("127.0.0.1", 0))
|
||||
self.srv.listen(1)
|
||||
self.port = self.srv.getsockname()[1]
|
||||
self._stop = False
|
||||
|
||||
def stop(self):
|
||||
self._stop = True
|
||||
try:
|
||||
self.srv.close()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def _last_begin(self, index):
|
||||
plen = self.meta.piece_len(index)
|
||||
nblocks = (plen + BLOCK - 1) // BLOCK
|
||||
return (nblocks - 1) * BLOCK
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
conn, _ = self.srv.accept()
|
||||
except OSError:
|
||||
return
|
||||
conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
hs = _recv_exact(conn, 68)
|
||||
if not hs:
|
||||
return
|
||||
info_hash = hs[28:48]
|
||||
conn.sendall(bytes([19]) + b"BitTorrent protocol" + bytes(8)
|
||||
+ info_hash + os.urandom(20))
|
||||
nbytes = (self.meta.num_pieces + 7) // 8
|
||||
bf = bytearray(nbytes)
|
||||
for i in range(self.meta.num_pieces):
|
||||
bf[i >> 3] |= 0x80 >> (i & 7)
|
||||
conn.sendall(_msg(5, bytes(bf))) # bitfield: has everything
|
||||
conn.sendall(_msg(1)) # unchoke
|
||||
|
||||
plen = self.meta.piece_length
|
||||
import struct
|
||||
while not self._stop:
|
||||
hdr = _recv_exact(conn, 4)
|
||||
if hdr is None:
|
||||
break
|
||||
ln = struct.unpack(">I", hdr)[0]
|
||||
if ln == 0:
|
||||
continue
|
||||
mid = _recv_exact(conn, 1)
|
||||
if mid is None:
|
||||
break
|
||||
payload = _recv_exact(conn, ln - 1) if ln > 1 else b""
|
||||
if payload is None:
|
||||
break
|
||||
if mid[0] != 6: # only act on REQUEST
|
||||
continue
|
||||
index, begin, length = struct.unpack(">III", payload)
|
||||
if begin == self._last_begin(index):
|
||||
continue # withhold the last block forever -> piece never finishes
|
||||
off = index * plen + begin
|
||||
conn.sendall(_piece_msg(index, begin, self.data[off:off + length]))
|
||||
self.served += 1
|
||||
|
||||
|
||||
class HealthySeed(threading.Thread):
|
||||
"""Plain seed that serves every requested block."""
|
||||
|
||||
def __init__(self, data, meta):
|
||||
super().__init__(daemon=True)
|
||||
self.data = data
|
||||
self.meta = meta
|
||||
self.srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
self.srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
self.srv.bind(("127.0.0.1", 0))
|
||||
self.srv.listen(1)
|
||||
self.port = self.srv.getsockname()[1]
|
||||
self._stop = False
|
||||
|
||||
def stop(self):
|
||||
self._stop = True
|
||||
try:
|
||||
self.srv.close()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
conn, _ = self.srv.accept()
|
||||
except OSError:
|
||||
return
|
||||
conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
hs = _recv_exact(conn, 68)
|
||||
if not hs:
|
||||
return
|
||||
info_hash = hs[28:48]
|
||||
conn.sendall(bytes([19]) + b"BitTorrent protocol" + bytes(8)
|
||||
+ info_hash + os.urandom(20))
|
||||
nbytes = (self.meta.num_pieces + 7) // 8
|
||||
bf = bytearray(nbytes)
|
||||
for i in range(self.meta.num_pieces):
|
||||
bf[i >> 3] |= 0x80 >> (i & 7)
|
||||
conn.sendall(_msg(5, bytes(bf)))
|
||||
conn.sendall(_msg(1))
|
||||
|
||||
plen = self.meta.piece_length
|
||||
import struct
|
||||
while not self._stop:
|
||||
hdr = _recv_exact(conn, 4)
|
||||
if hdr is None:
|
||||
break
|
||||
ln = struct.unpack(">I", hdr)[0]
|
||||
if ln == 0:
|
||||
continue
|
||||
mid = _recv_exact(conn, 1)
|
||||
if mid is None:
|
||||
break
|
||||
payload = _recv_exact(conn, ln - 1) if ln > 1 else b""
|
||||
if payload is None:
|
||||
break
|
||||
if mid[0] != 6:
|
||||
continue
|
||||
index, begin, length = struct.unpack(">III", payload)
|
||||
off = index * plen + begin
|
||||
conn.sendall(_piece_msg(index, begin, self.data[off:off + length]))
|
||||
|
||||
|
||||
def test_endgame_rescues_stalled_claims():
|
||||
ensure_built()
|
||||
with tempfile.TemporaryDirectory() as root:
|
||||
torrent, original, _ = make_torrent(root, 2 * 1024 * 1024) # 8 pieces
|
||||
meta = load_metadata(torrent)
|
||||
stalling = StallingSeed(original, meta)
|
||||
healthy = HealthySeed(original, meta)
|
||||
stalling.start()
|
||||
healthy.start()
|
||||
try:
|
||||
with Engine(EngineConfig(loop_count=1, slots_per_loop=1024,
|
||||
max_pipeline=64,
|
||||
request_timeout_ms=400), lib_path=LIB) as eng:
|
||||
tid = eng.add_torrent(meta.info_hash, make_peer_id(),
|
||||
meta.piece_length, meta.total_size,
|
||||
meta.num_pieces)
|
||||
eng.set_priorities(tid, [1] * meta.num_pieces)
|
||||
|
||||
# Let the stalling seed connect and claim pieces first.
|
||||
eng.add_peer(tid, "127.0.0.1", stalling.port)
|
||||
deadline = time.time() + 10.0
|
||||
while time.time() < deadline:
|
||||
if eng.status(tid).outstanding > 0:
|
||||
break
|
||||
eng.wait(20)
|
||||
else:
|
||||
raise TimeoutError("stalling seed never issued requests")
|
||||
|
||||
# Now the only way to finish the pieces it claimed is endgame.
|
||||
eng.add_peer(tid, "127.0.0.1", healthy.port)
|
||||
drives = {tid: TorrentDrive(meta)}
|
||||
drive_engine(eng, drives, timeout=30.0)
|
||||
|
||||
got = b"".join(bytes(b) for b in drives[tid].buffers)
|
||||
assert got == original, "endgame download bytes differ"
|
||||
finally:
|
||||
stalling.stop()
|
||||
healthy.stop()
|
||||
print("endgame OK: healthy seed finished pieces a stalled peer had claimed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_endgame_rescues_stalled_claims()
|
||||
Loading…
Add table
Add a link
Reference in a new issue