mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-14 13:05:09 +00:00
fix(serve): lock shared session state and throttle progress writes
Handlers share one service instance, so a per-session lock keeps two requests out of it at once. Track threads merge and write job state under one lock, and the progress sink coalesces ticks that say the same thing.
This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
"""A multi-period DASH manifest must fetch and probe the init segment of the first content
|
||||
period only.
|
||||
|
||||
``DASH.download_track`` keeps the init data and the KID of period 0 and discards them for
|
||||
every later period, so a request and an FFprobe process per extra period gain nothing. The
|
||||
``SegmentBase`` form is the exception: its ``media_range`` is built from ``len(init_data)``,
|
||||
so the request stays for every period and only the probe is dropped.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from functools import partial
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from unshackle.core.constants import DOWNLOAD_LICENCE_ONLY
|
||||
from unshackle.core.manifests import DASH
|
||||
from unshackle.core.tracks.track import DownloadContext, Track
|
||||
from unshackle.core.utils.xml import load_xml
|
||||
|
||||
MEDIA = bytes(i % 256 for i in range(4096))
|
||||
INIT_LEN = 100
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
server = self.server
|
||||
rng = self.headers.get("Range")
|
||||
with server.lock:
|
||||
server.requests.append((self.path, rng))
|
||||
body = server.routes.get(self.path)
|
||||
if body is None:
|
||||
self.send_response(404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
if rng:
|
||||
start, end = (int(x) for x in rng.removeprefix("bytes=").split("-"))
|
||||
chunk = body[start : end + 1]
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(body)}")
|
||||
body = chunk
|
||||
else:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||
srv.lock = threading.Lock()
|
||||
srv.requests = []
|
||||
srv.routes = {f"/p{i}/init.bin": MEDIA[:INIT_LEN] for i in (0, 1)}
|
||||
srv.routes.update({f"/p{i}/1.bin": MEDIA for i in (0, 1)})
|
||||
srv.routes.update({f"/p{i}.mp4": MEDIA for i in (0, 1)})
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def probes(monkeypatch):
|
||||
"""Count KID probes without running FFprobe."""
|
||||
calls = []
|
||||
|
||||
def counted(self, init_data=None, *args, **kwargs):
|
||||
calls.append(init_data)
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(Track, "get_key_id", counted)
|
||||
return calls
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def licence_only():
|
||||
"""Return from download_track once the periods are parsed, before any segment download."""
|
||||
was_set = DOWNLOAD_LICENCE_ONLY.is_set()
|
||||
DOWNLOAD_LICENCE_ONLY.set()
|
||||
yield
|
||||
if not was_set:
|
||||
DOWNLOAD_LICENCE_ONLY.clear()
|
||||
|
||||
|
||||
def base(srv) -> str:
|
||||
return f"http://127.0.0.1:{srv.server_address[1]}/"
|
||||
|
||||
|
||||
def template_mpd(srv) -> str:
|
||||
periods = "\n".join(
|
||||
f""" <Period id="{i}" duration="PT1S">
|
||||
<AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
|
||||
<Representation id="v0" codecs="avc1.640028" bandwidth="1000000" width="1920" height="1080">
|
||||
<SegmentTemplate timescale="1" duration="1" startNumber="1"
|
||||
initialization="p{i}/init.bin" media="p{i}/$Number$.bin"/>
|
||||
</Representation>
|
||||
</AdaptationSet>
|
||||
</Period>"""
|
||||
for i in (0, 1)
|
||||
)
|
||||
return f"""<?xml version="1.0"?>
|
||||
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT2S">
|
||||
<BaseURL>{base(srv)}</BaseURL>
|
||||
{periods}
|
||||
</MPD>"""
|
||||
|
||||
|
||||
def base_mpd(srv) -> str:
|
||||
periods = "\n".join(
|
||||
f""" <Period id="{i}" duration="PT1S">
|
||||
<AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
|
||||
<Representation id="v0" codecs="avc1.640028" bandwidth="1000000" width="1920" height="1080">
|
||||
<BaseURL>p{i}.mp4</BaseURL>
|
||||
<SegmentBase>
|
||||
<Initialization range="0-{INIT_LEN - 1}"/>
|
||||
</SegmentBase>
|
||||
</Representation>
|
||||
</AdaptationSet>
|
||||
</Period>"""
|
||||
for i in (0, 1)
|
||||
)
|
||||
return f"""<?xml version="1.0"?>
|
||||
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" mediaPresentationDuration="PT2S">
|
||||
<BaseURL>{base(srv)}</BaseURL>
|
||||
{periods}
|
||||
</MPD>"""
|
||||
|
||||
|
||||
def make_track(mpd: str, srv):
|
||||
return DASH.from_text(mpd, url=f"{base(srv)}manifest.mpd").to_tracks(language="en").videos[0]
|
||||
|
||||
|
||||
def make_ctx(tmp_path):
|
||||
return DownloadContext(
|
||||
save_path=tmp_path / "out.mp4",
|
||||
save_dir=tmp_path / "v0_segments",
|
||||
progress=partial(lambda **kw: None),
|
||||
)
|
||||
|
||||
|
||||
def test_segment_template_init_is_fetched_once_for_two_periods(server, probes, licence_only, tmp_path):
|
||||
track = make_track(template_mpd(server), server)
|
||||
|
||||
DASH.download_track(track, make_ctx(tmp_path))
|
||||
|
||||
# both periods contributed a segment, so period 1 was parsed, not skipped
|
||||
assert len(track.data["dash"]["segment_durations"]) == 2
|
||||
assert [path for path, _ in server.requests] == ["/p0/init.bin"]
|
||||
assert len(probes) == 1
|
||||
|
||||
|
||||
def test_segment_base_keeps_each_request_but_probes_once(server, probes, licence_only, tmp_path):
|
||||
track = make_track(base_mpd(server), server)
|
||||
|
||||
DASH.download_track(track, make_ctx(tmp_path))
|
||||
|
||||
# media_range needs the init bytes of every period, so every period keeps its request
|
||||
assert [path for path, _ in server.requests] == ["/p0.mp4", "/p1.mp4"]
|
||||
assert len(probes) == 1
|
||||
|
||||
|
||||
def test_segment_base_range_survives_a_skipped_probe(server, probes):
|
||||
tree = load_xml(base_mpd(server).encode())
|
||||
period = tree.findall("Period")[1]
|
||||
aset = period.find("AdaptationSet")
|
||||
rep = aset.find("Representation")
|
||||
|
||||
_, segments, _, _, track_kid = DASH.get_period_segments(
|
||||
period,
|
||||
aset,
|
||||
rep,
|
||||
tree,
|
||||
Track(url=base(server), language="en"),
|
||||
f"{base(server)}manifest.mpd",
|
||||
requests.Session(),
|
||||
probe_kid=False,
|
||||
)
|
||||
|
||||
assert segments == [(f"{base(server)}p1.mp4", f"{INIT_LEN}-{len(MEDIA)}")]
|
||||
assert track_kid is None
|
||||
assert not probes
|
||||
@@ -0,0 +1,143 @@
|
||||
"""File I/O rules for HLS segments.
|
||||
|
||||
Two faults that Linux never shows:
|
||||
|
||||
* The FFmpeg concat list is a UTF-8 file by specification. Written in the OS locale
|
||||
encoding it either raises ``UnicodeEncodeError`` on a non-ASCII path, which kills the
|
||||
track after every segment is already on disk, or it silently disables the fast path.
|
||||
* A subtitle segment that the UTF-8 transform leaves unchanged, the common case, needs
|
||||
no write back. Each rewrite is a file operation an on-access virus scanner has to
|
||||
clear.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import codecs
|
||||
import subprocess
|
||||
import threading
|
||||
from functools import partial
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from unshackle.core.manifests import HLS
|
||||
from unshackle.core.tracks.track import DownloadContext
|
||||
|
||||
SEGMENT_ONE = b"WEBVTT\n\n00:00:00.000 --> 00:00:04.000\nhello\n"
|
||||
SEGMENT_TWO = b"WEBVTT\n\n00:00:04.000 --> 00:00:08.000\nworld\n"
|
||||
|
||||
|
||||
def test_concat_list_is_written_as_utf8(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A non-ASCII segment path must not depend on the OS codepage."""
|
||||
save_dir = tmp_path / "Ünshackle тест" / "Subtitle_segments"
|
||||
save_dir.mkdir(parents=True)
|
||||
segments = []
|
||||
for i in range(2):
|
||||
segment = save_dir / f"{i:05}.vtt"
|
||||
segment.write_bytes(b"\x00" * 8)
|
||||
segments.append(segment)
|
||||
|
||||
encodings: list[object] = []
|
||||
real_write_text = Path.write_text
|
||||
|
||||
def spy_write_text(self, data, encoding=None, *args, **kwargs):
|
||||
encodings.append(encoding)
|
||||
return real_write_text(self, data, encoding=encoding, *args, **kwargs)
|
||||
|
||||
seen: list[bytes] = []
|
||||
|
||||
def fake_run(args, **kwargs):
|
||||
seen.append(Path(args[args.index("-i") + 1]).read_bytes())
|
||||
raise subprocess.CalledProcessError(1, args, output="", stderr="")
|
||||
|
||||
monkeypatch.setattr(Path, "write_text", spy_write_text)
|
||||
monkeypatch.setattr("unshackle.core.manifests.hls.binaries.FFMPEG", Path("ffmpeg"))
|
||||
monkeypatch.setattr("unshackle.core.manifests.hls.subprocess.run", fake_run)
|
||||
|
||||
HLS.merge_segments(segments=segments, save_path=save_dir.parent / "Subtitle.vtt")
|
||||
|
||||
assert encodings, "the concat list was never written"
|
||||
assert all(e is not None and codecs.lookup(e).name == "utf-8" for e in encodings), (
|
||||
f"the concat list was written in the OS locale encoding: {encodings}"
|
||||
)
|
||||
assert seen and seen[0].decode("utf-8") == "\n".join(f"file '{s.absolute()}'" for s in segments)
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
body = self.server.routes.get(self.path)
|
||||
if body is None:
|
||||
self.send_response(404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||
srv.routes = {}
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
MASTER = """#EXTM3U
|
||||
#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID="subs",NAME="English",LANGUAGE="en",DEFAULT=YES,URI="subs.m3u8"
|
||||
#EXT-X-STREAM-INF:BANDWIDTH=1000000,RESOLUTION=1920x1080,CODECS="avc1.640028",SUBTITLES="subs"
|
||||
media.m3u8
|
||||
"""
|
||||
|
||||
SUBS_PLAYLIST = b"""#EXTM3U
|
||||
#EXT-X-VERSION:3
|
||||
#EXT-X-TARGETDURATION:4
|
||||
#EXTINF:4.0,
|
||||
part0.vtt
|
||||
#EXTINF:4.0,
|
||||
part1.vtt
|
||||
#EXT-X-ENDLIST
|
||||
"""
|
||||
|
||||
|
||||
def test_unchanged_subtitle_segments_are_not_rewritten(server, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Plain UTF-8 WebVTT survives the transform untouched, so nothing needs writing back."""
|
||||
port = server.server_address[1]
|
||||
server.routes["/subs.m3u8"] = SUBS_PLAYLIST
|
||||
server.routes["/part0.vtt"] = SEGMENT_ONE
|
||||
server.routes["/part1.vtt"] = SEGMENT_TWO
|
||||
|
||||
track = HLS.from_text(MASTER, url=f"http://127.0.0.1:{port}/master.m3u8").to_tracks(language="en").subtitles[0]
|
||||
|
||||
save_dir = tmp_path / "s0_segments"
|
||||
rewritten: list[str] = []
|
||||
real_write_bytes = Path.write_bytes
|
||||
|
||||
def spy_write_bytes(self, data):
|
||||
if save_dir in self.parents:
|
||||
rewritten.append(self.name)
|
||||
return real_write_bytes(self, data)
|
||||
|
||||
monkeypatch.setattr(Path, "write_bytes", spy_write_bytes)
|
||||
|
||||
HLS.download_track(
|
||||
track,
|
||||
DownloadContext(
|
||||
save_path=tmp_path / "out.vtt",
|
||||
save_dir=save_dir,
|
||||
progress=partial(lambda **kw: None),
|
||||
max_workers=1,
|
||||
),
|
||||
)
|
||||
|
||||
assert (tmp_path / "out.vtt").read_bytes() == SEGMENT_ONE + SEGMENT_TWO
|
||||
assert rewritten == [], f"unchanged subtitle segments were rewritten: {rewritten}"
|
||||
@@ -0,0 +1,65 @@
|
||||
"""A reused CustomRemoteCDM must not carry one title's PlayReady PSSH into the next.
|
||||
|
||||
``dl.get_cdm`` builds each CDM once per run, so one instance serves every title.
|
||||
``PlayReady.get_content_keys`` calls ``set_pssh_b64`` only when the DRM has a PSSH, so
|
||||
an instance-level value left over from an earlier title would be sent as the init data
|
||||
of a title that has none.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from unshackle.core.cdm.custom_remote_cdm import CustomRemoteCDM
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
class FakePssh:
|
||||
"""A PSSH object shaped like the ones the manifests hand to the CDM."""
|
||||
|
||||
def __init__(self, raw: bytes) -> None:
|
||||
self.raw = raw
|
||||
|
||||
|
||||
def make_cdm() -> CustomRemoteCDM:
|
||||
return CustomRemoteCDM(host="https://cdm.test", device={"name": "SL3000", "type": "PLAYREADY"})
|
||||
|
||||
|
||||
def set_pssh(cdm: CustomRemoteCDM, session_id: bytes, pssh_b64: str) -> None:
|
||||
"""Store a PSSH the way ``PlayReady.get_content_keys`` does."""
|
||||
try:
|
||||
cdm.set_pssh_b64(pssh_b64, session_id=session_id)
|
||||
except TypeError: # CDM predating the per-session signature
|
||||
cdm.set_pssh_b64(pssh_b64)
|
||||
|
||||
|
||||
def test_a_later_session_does_not_inherit_the_earlier_pssh() -> None:
|
||||
cdm = make_cdm()
|
||||
assert cdm.is_playready
|
||||
|
||||
first = cdm.open()
|
||||
set_pssh(cdm, first, "FIRST_TITLE_PSSH")
|
||||
assert cdm.get_init_data_from_pssh(FakePssh(b"ignored")) == "FIRST_TITLE_PSSH"
|
||||
cdm.close(first)
|
||||
|
||||
# The second title has no PSSH of its own, so nothing calls set_pssh_b64.
|
||||
cdm.open()
|
||||
init_data = cdm.get_init_data_from_pssh(FakePssh(b"second"))
|
||||
|
||||
assert init_data == base64.b64encode(b"second").decode("utf-8")
|
||||
|
||||
|
||||
def test_two_open_sessions_keep_their_own_pssh() -> None:
|
||||
"""Tracks decrypt in parallel on the one shared instance, so the store is per session."""
|
||||
cdm = make_cdm()
|
||||
|
||||
one = cdm.open()
|
||||
two = cdm.open()
|
||||
set_pssh(cdm, one, "PSSH_ONE")
|
||||
set_pssh(cdm, two, "PSSH_TWO")
|
||||
|
||||
assert cdm.get_init_data_from_pssh(FakePssh(b"x"), cdm._sessions[one]["pssh_b64"]) == "PSSH_ONE"
|
||||
assert cdm.get_init_data_from_pssh(FakePssh(b"x"), cdm._sessions[two]["pssh_b64"]) == "PSSH_TWO"
|
||||
@@ -0,0 +1,110 @@
|
||||
"""A bounded sub-range that returns no bytes must fail, not spin.
|
||||
|
||||
After a cut, the sequential path asks for bounded sub-ranges and treats a whole
|
||||
sub-range as progress rather than a retry. A response with an empty body of
|
||||
unknown length (a chunked 206, a 204, a 304) makes no progress, so it must
|
||||
raise and spend an attempt, or the same request repeats forever.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
import requests as rq
|
||||
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED
|
||||
from unshackle.core.downloaders.requests import download
|
||||
|
||||
# the package re-exports the `requests` function, shadowing the module of the same name
|
||||
downloader = sys.modules[download.__module__]
|
||||
|
||||
BODY = bytes(range(256)) * 256 # 64 KiB
|
||||
CUT_AFTER = 10 * 1024
|
||||
CHUNK = 16 * 1024
|
||||
DEADLINE = 30
|
||||
|
||||
|
||||
class _EmptyRangeServer(ThreadingHTTPServer):
|
||||
"""Cuts the first body, then answers every bounded range with an empty chunked 206."""
|
||||
|
||||
allow_reuse_address = True
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self):
|
||||
self.ranges: list[str] = []
|
||||
super().__init__(("127.0.0.1", 0), _EmptyRangeHandler)
|
||||
self.url = f"http://127.0.0.1:{self.server_address[1]}/seg.mp4"
|
||||
|
||||
def handle_error(self, request, client_address): # the client hanging up is expected here
|
||||
pass
|
||||
|
||||
|
||||
class _EmptyRangeHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def log_message(self, *_):
|
||||
pass
|
||||
|
||||
def do_GET(self):
|
||||
header = self.headers.get("Range")
|
||||
self.server.ranges.append(header)
|
||||
start, end, bounded = 0, len(BODY) - 1, False
|
||||
if header:
|
||||
start_s, _, end_s = header.removeprefix("bytes=").partition("-")
|
||||
start, end = int(start_s), (int(end_s) if end_s else len(BODY) - 1)
|
||||
bounded = bool(end_s)
|
||||
if bounded:
|
||||
# a valid response: the chunked terminator closes a zero-length body, so requests raises nothing
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(BODY)}")
|
||||
self.send_header("Transfer-Encoding", "chunked")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
return
|
||||
body = BODY[start : end + 1]
|
||||
self.send_response(206 if header else 200)
|
||||
if header:
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(BODY)}")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
# the cut teaches the downloader to bound the requests that follow
|
||||
self.wfile.write(body[:CUT_AFTER])
|
||||
self.close_connection = True
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server(monkeypatch):
|
||||
DOWNLOAD_CANCELLED.clear()
|
||||
monkeypatch.setattr(downloader, "RETRY_WAIT", 0)
|
||||
monkeypatch.setattr(downloader, "REQUEST_CHUNK_SIZE", CHUNK)
|
||||
monkeypatch.setattr(downloader, "MIN_REQUEST_CHUNK", 4 * 1024)
|
||||
monkeypatch.setattr(downloader, "MIN_RESUME_PROGRESS", 1024)
|
||||
server = _EmptyRangeServer()
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
yield server
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
DOWNLOAD_CANCELLED.clear()
|
||||
|
||||
|
||||
def test_empty_sub_range_fails_instead_of_looping(tmp_path, server):
|
||||
save_path = tmp_path / "0.mp4"
|
||||
result: list[BaseException] = []
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
list(download(url=server.url, save_path=save_path, session=rq.Session()))
|
||||
except BaseException as exc:
|
||||
result.append(exc)
|
||||
|
||||
# a hard deadline: a regression hangs this thread, and the suite must fail, not stall
|
||||
worker = threading.Thread(target=run, daemon=True)
|
||||
worker.start()
|
||||
worker.join(timeout=DEADLINE)
|
||||
|
||||
assert not worker.is_alive(), f"download still running after {DEADLINE}s ({len(server.ranges)} requests)"
|
||||
assert result and isinstance(result[0], IOError), f"expected an IOError, got {result!r}"
|
||||
assert "no bytes" in str(result[0])
|
||||
assert len(server.ranges) <= downloader.MAX_ATTEMPTS + 1
|
||||
assert not save_path.exists()
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Connection pooling survives the no-retry adapter view.
|
||||
|
||||
``download()`` builds this view once per segment. ``copy()`` runs
|
||||
``HTTPAdapter.__setstate__``, which blanks ``proxy_manager`` and re-runs
|
||||
``init_poolmanager()``, so a view that keeps those fresh objects hands every segment an
|
||||
empty pool: one TCP connection and one TLS handshake each. That is invisible on Linux and
|
||||
costs a native Windows host roughly 7x the wall time and 4x the CPU on a manifest of a few
|
||||
thousand segments.
|
||||
|
||||
The second test pins the reason sharing one pool between the two adapters is safe:
|
||||
``HTTPAdapter.send`` passes ``retries=self.max_retries`` per call, so the view's ``Retry(0)``
|
||||
travels with the view and never reaches the caller's own session.
|
||||
"""
|
||||
|
||||
from requests import Session
|
||||
from requests.adapters import HTTPAdapter
|
||||
from urllib3.util.retry import Retry
|
||||
|
||||
from unshackle.core.downloaders.requests import no_retry_session
|
||||
|
||||
|
||||
def pooled_session(total: int = 5) -> Session:
|
||||
session = Session()
|
||||
adapter = HTTPAdapter(pool_connections=16, pool_maxsize=16, pool_block=True, max_retries=Retry(total=total))
|
||||
session.mount("https://", adapter)
|
||||
session.mount("http://", adapter)
|
||||
return session
|
||||
|
||||
|
||||
def test_view_shares_the_pools_it_was_copied_from() -> None:
|
||||
session = pooled_session()
|
||||
view = no_retry_session(session)
|
||||
|
||||
for prefix, adapter in session.adapters.items():
|
||||
twin = view.adapters[prefix]
|
||||
assert twin is not adapter, "the view must be its own adapter, or Retry(0) lands on the caller"
|
||||
assert twin.poolmanager is adapter.poolmanager
|
||||
assert twin.proxy_manager is adapter.proxy_manager
|
||||
|
||||
|
||||
def test_view_never_retries_and_leaves_the_caller_retrying() -> None:
|
||||
session = pooled_session(total=5)
|
||||
view = no_retry_session(session)
|
||||
|
||||
assert view.adapters["https://"].max_retries.total == 0
|
||||
assert session.adapters["https://"].max_retries.total == 5, "the caller's retry policy must survive"
|
||||
|
||||
|
||||
def test_repeated_views_keep_reusing_one_pool() -> None:
|
||||
"""``download()`` rebuilds the view per segment; every rebuild must land on the same pool."""
|
||||
session = pooled_session()
|
||||
origin = session.adapters["https://"].poolmanager
|
||||
|
||||
for _ in range(3):
|
||||
assert no_retry_session(session).adapters["https://"].poolmanager is origin
|
||||
@@ -18,7 +18,7 @@ def load_scenarios(config: pytest.Config) -> dict:
|
||||
import yaml
|
||||
|
||||
path = Path(__file__).parent / "fixtures" / "fixtures.yaml"
|
||||
data = yaml.safe_load(path.read_text()) if path.exists() else {}
|
||||
data = yaml.safe_load(path.read_text(encoding="utf-8")) if path.exists() else {}
|
||||
selected = (config.getoption("--services") or "").strip()
|
||||
services = (data or {}).get("services", {}) or {}
|
||||
if selected:
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Blocking service calls must not run on the aiohttp event loop.
|
||||
|
||||
get_titles and get_tracks do network work that can take a whole manifest round trip. Run on
|
||||
the loop, they stop the server: the progress pump and the SSE keep-alive stall with them. An
|
||||
interactive prompt from the service can never be answered either, because the endpoint that
|
||||
answers it needs the same loop. Each test here holds the service call in a two-event
|
||||
handshake that only completes while the loop still runs another coroutine.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Awaitable, Callable, List
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from unshackle.core.api import handlers
|
||||
from unshackle.core.api.session_store import get_session_store
|
||||
from unshackle.core.titles.movie import Movie
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
# Bounds the wait when the loop is blocked, so a regression fails the test instead of hanging.
|
||||
HANDSHAKE_TIMEOUT = 5.0
|
||||
|
||||
|
||||
class Handshake:
|
||||
"""A blocking call that reports whether the event loop kept running during it."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.entered = threading.Event()
|
||||
self.resume = threading.Event()
|
||||
self.loop_was_blocked = False
|
||||
|
||||
def block(self) -> None:
|
||||
self.entered.set()
|
||||
self.loop_was_blocked = not self.resume.wait(timeout=HANDSHAKE_TIMEOUT)
|
||||
|
||||
async def drive(self, call: Callable[[], Awaitable[web.Response]]) -> web.Response:
|
||||
task = asyncio.create_task(call())
|
||||
# Reached only while the loop is free.
|
||||
await asyncio.to_thread(self.entered.wait, HANDSHAKE_TIMEOUT)
|
||||
self.resume.set()
|
||||
return await asyncio.wait_for(task, HANDSHAKE_TIMEOUT)
|
||||
|
||||
|
||||
def a_movie() -> Movie:
|
||||
return Movie(id_="movie-0001", service=object, name="Film", year=2024, language=None)
|
||||
|
||||
|
||||
class BlockingTitles:
|
||||
def __init__(self, handshake: Handshake) -> None:
|
||||
self.handshake = handshake
|
||||
|
||||
def get_titles(self) -> List[Movie]:
|
||||
self.handshake.block()
|
||||
return [a_movie()]
|
||||
|
||||
|
||||
class BlockingTracks:
|
||||
def __init__(self, handshake: Handshake) -> None:
|
||||
self.handshake = handshake
|
||||
|
||||
def get_titles(self) -> List[Movie]:
|
||||
return [a_movie()]
|
||||
|
||||
def get_tracks(self, title: Any) -> Any:
|
||||
self.handshake.block()
|
||||
return SimpleNamespace(videos=[], audio=[], subtitles=[])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_titles_leaves_the_loop_free() -> None:
|
||||
handshake = Handshake()
|
||||
store = get_session_store()
|
||||
await store.create("EXAMPLE", BlockingTitles(handshake), session_id="loop-test-titles")
|
||||
try:
|
||||
response = await handshake.drive(lambda: handlers.session_titles_handler("loop-test-titles"))
|
||||
finally:
|
||||
await store.delete("loop-test-titles")
|
||||
|
||||
assert not handshake.loop_was_blocked, "get_titles blocked the event loop"
|
||||
assert [t["name"] for t in json.loads(response.body)["titles"]] == ["Film"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tracks_leaves_the_loop_free(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
handshake = Handshake()
|
||||
service = BlockingTracks(handshake)
|
||||
monkeypatch.setattr(handlers, "validate_service", lambda tag, request=None: "EXAMPLE")
|
||||
monkeypatch.setattr(handlers, "setup_list_service", lambda *args, **kwargs: service)
|
||||
|
||||
data = {"service": "EXAMPLE", "title_id": "movie-0001"}
|
||||
response = await handshake.drive(lambda: handlers.list_tracks_handler(data))
|
||||
|
||||
assert not handshake.loop_was_blocked, "get_tracks blocked the event loop"
|
||||
assert json.loads(response.body)["title"]["name"] == "Film"
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Unit tests for the progress write rate and for concurrent progress writers.
|
||||
|
||||
The aggregate sink rewrites a whole JSON file per call, and one track sends thousands of ticks
|
||||
while the reader polls twice a second. These tests pin the two guards on that: repeat ticks are
|
||||
throttled, a tick that changes the job shape is not, and two threads writing the same file at the
|
||||
same time do not lose a write."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from unshackle.core.api import progress as progress_module
|
||||
from unshackle.core.api.download_worker import write_result
|
||||
from unshackle.core.api.progress import build_job_progress_callables
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
class Video:
|
||||
def __init__(self, height=1080, bitrate=4_000_000):
|
||||
self.height = height
|
||||
self.range = None
|
||||
self.bitrate = bitrate
|
||||
|
||||
|
||||
def noop(**kwargs):
|
||||
pass
|
||||
|
||||
|
||||
def test_repeat_ticks_are_throttled():
|
||||
updates: list[dict] = []
|
||||
cbs = build_job_progress_callables([Video()], [noop], updates.append)
|
||||
|
||||
for completed in range(1, 501):
|
||||
cbs[0](total=1000, completed=completed)
|
||||
|
||||
# One emit for the first tick (the track starts); the rest say the same thing and coalesce.
|
||||
assert len(updates) == 1
|
||||
|
||||
|
||||
def test_track_finish_bypasses_the_throttle():
|
||||
updates: list[dict] = []
|
||||
cbs = build_job_progress_callables([Video(), Video()], [noop, noop], updates.append)
|
||||
|
||||
cbs[0](total=100, completed=1)
|
||||
cbs[0](total=100, completed=2)
|
||||
cbs[0](total=100, completed=3)
|
||||
assert len(updates) == 1
|
||||
|
||||
cbs[0](downloaded="Downloaded")
|
||||
assert len(updates) == 2
|
||||
assert updates[-1]["completed_tracks"] == 1
|
||||
|
||||
|
||||
def test_phase_change_bypasses_the_throttle():
|
||||
updates: list[dict] = []
|
||||
cbs = build_job_progress_callables([Video(2160), Video(720)], [noop, noop], updates.append)
|
||||
|
||||
cbs[0](total=100, completed=1)
|
||||
first_phase = updates[-1]["phase"]
|
||||
|
||||
cbs[1](total=100, completed=1)
|
||||
assert len(updates) == 2
|
||||
assert updates[-1]["phase"] != first_phase
|
||||
|
||||
|
||||
def test_filled_track_bypasses_the_throttle_without_a_terminal_string():
|
||||
"""A track can reach its segment total and never send "Downloaded". That last tick must land,
|
||||
or the job reports a percentage that is short for good."""
|
||||
updates: list[dict] = []
|
||||
cbs = build_job_progress_callables([Video()], [noop], updates.append)
|
||||
|
||||
cbs[0](total=100, completed=10)
|
||||
cbs[0](total=100, completed=50)
|
||||
assert len(updates) == 1
|
||||
|
||||
cbs[0](total=100, completed=100)
|
||||
assert len(updates) == 2
|
||||
assert updates[-1]["progress"] == pytest.approx(progress_module.DOWNLOAD_PROGRESS_CEILING)
|
||||
|
||||
|
||||
def test_throttle_releases_after_the_interval(monkeypatch: pytest.MonkeyPatch):
|
||||
clock = {"now": 1000.0}
|
||||
monkeypatch.setattr(progress_module.time, "monotonic", lambda: clock["now"])
|
||||
|
||||
updates: list[dict] = []
|
||||
cbs = build_job_progress_callables([Video()], [noop], updates.append)
|
||||
|
||||
cbs[0](total=100, completed=1)
|
||||
cbs[0](total=100, completed=2)
|
||||
assert len(updates) == 1
|
||||
|
||||
clock["now"] += progress_module.PROGRESS_EMIT_INTERVAL + 0.01
|
||||
cbs[0](total=100, completed=3)
|
||||
assert len(updates) == 2
|
||||
|
||||
|
||||
def test_concurrent_writers_do_not_clobber_each_other(tmp_path: Path):
|
||||
"""Two threads can write the same progress file. A temp name shared by two writers lets one
|
||||
replace away the file the other is about to move, which fails that write."""
|
||||
target = tmp_path / "progress.json"
|
||||
errors: list[BaseException] = []
|
||||
start = threading.Barrier(8)
|
||||
|
||||
def writer(index: int) -> None:
|
||||
start.wait()
|
||||
try:
|
||||
for tick in range(40):
|
||||
write_result(target, {"writer": index, "tick": tick})
|
||||
except BaseException as exc: # noqa: BLE001 - reported to the test thread
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=writer, args=(i,)) for i in range(8)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
assert errors == []
|
||||
payload = json.loads(target.read_text(encoding="utf-8"))
|
||||
assert set(payload) == {"writer", "tick"}
|
||||
assert [p.name for p in tmp_path.iterdir()] == ["progress.json"]
|
||||
Reference in New Issue
Block a user