Naut-Peer/tests/test_endgame.py

"""
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()