unshackle/tests/downloaders/test_probe_ranged.py

"""``probe_ranged`` cost and fidelity.

The probe asks for one byte to learn the complete length and whether the host honours
ranges. It runs once for every single-URL download and once per tail-boost candidate,
so what it costs and what it sends both matter:

- it must leave the connection reusable. A streamed response whose body is never read
  cannot go back in the pool. Every probe would then tear down a socket, and the next
  request would pay a fresh handshake.
- it must ask for an unencoded body like ``download()`` does. A CDN that answers a
  probe with a Content-Encoding makes the probe decline, and the ranged-parallel path
  is silently lost for a file that supports it.
- it must carry the URL item's own request options. Without them a host that needs an
  auth header answers the probe with a 4xx, and the download falls back to one
  connection even though the ranges were there.

The drain is bounded by the declared Content-Length, not by the status code. A host
that ignores Range may answer 200 with the whole resource, and one that answers 206
with the whole resource is just as possible. Reading either would download the file
only to decide not to use it.
"""

import importlib
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import pytest
from requests import Session

from unshackle.core.constants import DOWNLOAD_CANCELLED

# the package re-exports the `requests` function, shadowing the module of the same name
dl = importlib.import_module("unshackle.core.downloaders.requests")

PAYLOAD = bytes((i * 11 + 3) % 256 for i in range(32 * 1024))
BIG = bytes(8 * 1024 * 1024)  # larger than any socket buffer, so a full drain is visible
PROBES = 6
AUTH = "let-me-in"


class _ProbeServer(ThreadingHTTPServer):
    allow_reuse_address = True
    daemon_threads = True

    def __init__(self):
        self.lock = threading.Lock()
        self.requests: list = []  # (Range, Accept-Encoding, X-Auth)
        self.connections = 0
        self.mode = "range"
        self.written = 0
        self.broke = False
        super().__init__(("127.0.0.1", 0), _ProbeHandler)

    def handle_error(self, request, client_address):  # the client hanging up is expected here
        pass


class _ProbeHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *_):
        pass

    def setup(self):
        super().setup()
        with self.server.lock:
            self.server.connections += 1

    def do_GET(self):
        srv = self.server
        rng = self.headers.get("Range")
        with srv.lock:
            srv.requests.append((rng, self.headers.get("Accept-Encoding"), self.headers.get("X-Auth")))

        if srv.mode == "needs_auth" and self.headers.get("X-Auth") != AUTH:
            self.send_response(403)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        if srv.mode == "ignores_range":
            self.send_response(200)
            self.send_header("Content-Length", str(len(BIG)))
            self.end_headers()
            self._write_slowly(BIG)
            return

        if srv.mode == "fat_206":
            # a 206 status on a whole-resource body: the status alone cannot bound the drain
            self.send_response(206)
            self.send_header("Content-Range", f"bytes 0-{len(BIG) - 1}/{len(BIG)}")
            self.send_header("Content-Length", str(len(BIG)))
            self.end_headers()
            self._write_slowly(BIG)
            return

        if not rng:
            self.send_response(200)
            self.send_header("Content-Length", str(len(PAYLOAD)))
            self.end_headers()
            self.wfile.write(PAYLOAD)
            return

        start_s, _, end_s = rng.removeprefix("bytes=").partition("-")
        start = int(start_s)
        end = int(end_s) if end_s else len(PAYLOAD) - 1
        body = PAYLOAD[start : end + 1]
        self.send_response(206)
        self.send_header("Content-Range", f"bytes {start}-{end}/{len(PAYLOAD)}")
        self.send_header("Content-Length", str(len(body)))
        if srv.mode == "encode_unless_identity" and self.headers.get("Accept-Encoding") != "identity":
            self.send_header("Content-Encoding", "gzip")
        self.end_headers()
        self.wfile.write(body)

    def _write_slowly(self, body):
        # chunked writes so the count records how far the client actually read
        srv = self.server
        try:
            for at in range(0, len(body), 64 * 1024):
                self.wfile.write(body[at : at + 64 * 1024])
                with srv.lock:
                    srv.written += len(body[at : at + 64 * 1024])
        except OSError:
            with srv.lock:
                srv.broke = True
            self.close_connection = True


@pytest.fixture
def server(monkeypatch):
    DOWNLOAD_CANCELLED.clear()
    monkeypatch.setattr(dl, "RETRY_WAIT", 0)
    srv = _ProbeServer()
    thread = threading.Thread(target=srv.serve_forever, daemon=True)
    thread.start()
    yield srv
    srv.shutdown()
    thread.join(timeout=5)
    srv.server_close()
    DOWNLOAD_CANCELLED.clear()


def url(srv, path="/payload.bin"):
    return f"http://127.0.0.1:{srv.server_address[1]}{path}"


def test_probe_reuses_the_pooled_connection(server):
    session = Session()

    for _ in range(PROBES):
        assert dl.probe_ranged(url(server), session) == (len(PAYLOAD), True)

    assert server.connections == 1


def test_probe_does_not_drain_a_200_body(server):
    server.mode = "ignores_range"
    session = Session()

    assert dl.probe_ranged(url(server), session) == (0, False)

    for _ in range(50):
        with server.lock:
            if server.broke or server.written >= len(BIG):
                break
        time.sleep(0.1)
    with server.lock:
        assert server.written < len(BIG)


def test_probe_asks_for_identity_encoding(server):
    server.mode = "encode_unless_identity"
    session = Session()

    assert dl.probe_ranged(url(server), session) == (len(PAYLOAD), True)
    assert server.requests[-1][1] == "identity"

    # a caller that names its own encoding still wins
    assert dl.probe_ranged(url(server), session, headers={"Accept-Encoding": "gzip"}) == (0, False)
    assert server.requests[-1][1] == "gzip"


def test_probe_uses_the_url_items_headers(server, tmp_path, monkeypatch):
    server.mode = "needs_auth"
    monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024)
    item = {"url": url(server), "headers": {"X-Auth": AUTH}}

    for _ in dl.requests([item], output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4):
        pass

    save_path = tmp_path / "seg_0000.bin"
    assert save_path.read_bytes() == PAYLOAD
    probes = [r for r in server.requests if r[0] == "bytes=0-0"]
    assert probes and all(r[2] == AUTH for r in probes)
    # the probe was answered, so the ranged path ran instead of one sequential connection
    assert [r for r in server.requests if r[0] and r[0] != "bytes=0-0"]


def test_probe_does_not_drain_a_fat_206_body(server):
    """A 206 carrying the whole resource must not be read to decide whether to use ranges."""
    server.mode = "fat_206"
    session = Session()

    assert dl.probe_ranged(url(server), session) == (len(BIG), True)

    for _ in range(50):
        with server.lock:
            if server.broke or server.written >= len(BIG):
                break
        time.sleep(0.1)
    with server.lock:
        assert server.written < len(BIG)