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.

tests/core/test_dash_multi_period_init.py

@@ -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

tests/core/test_hls_segment_io.py

@@ -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}"

tests/core/test_remote_cdm_pssh_reuse.py

@@ -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"