From 79d109d548973538a1f8dbb2a5736623957e693d Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Mon, 21 Sep 2026 19:13:39 -0600 Subject: [PATCH] feat(dash): add merge_segments to append segments during the download Opt-in rolling merge for DASH video and audio: each contiguous segment is appended to the output as it lands, so the post-download "Merging" stage goes away. Decrypts the init segment eagerly when decrypt_segments is on, disables resume for the track, and skips subtitle tracks. --- docs/reference/configuration/download.md | 36 +++++ docs/reference/configuration/index.md | 2 +- tests/core/test_dash_collapse_integration.py | 31 ++++ tests/test_rolling_merge.py | 51 +++++++ unshackle/core/config.py | 1 + unshackle/core/drm/segment_decrypt.py | 25 +++- unshackle/core/manifests/dash.py | 146 +++++++++++++++---- unshackle/core/tracks/track.py | 4 +- 8 files changed, 257 insertions(+), 39 deletions(-) create mode 100644 tests/test_rolling_merge.py diff --git a/docs/reference/configuration/download.md b/docs/reference/configuration/download.md index c685214b..cc40f586 100644 --- a/docs/reference/configuration/download.md +++ b/docs/reference/configuration/download.md @@ -120,6 +120,9 @@ ranges must all match the previous download, or the download silently restarts c downloader has its own in-flight partial-file handling for those). - HLS resumes Widevine/PlayReady and unencrypted titles only. An AES-128 title always restarts. An HLS failure during the post-download merge step also restarts. + - A track downloaded with [`merge_segments`](#merge_segments) or + [`decrypt_segments`](drm.md#decrypt_segments) never resumes. Both options consume + or rewrite the segment files during the download. For a one-off resume without changing config, pass `dl --continue-downloads` on the command line (the flag can only enable resume, never disable it). @@ -128,6 +131,39 @@ command line (the flag can only enable resume, never disable it). continue_downloads: true ``` +## `merge_segments` + +- **Type:** `bool`  ·  **Default:** `false` + +Appends each DASH segment to the output file as soon as it and every segment before it have +arrived, instead of joining all the segment files after the last one lands. The join runs +during the download, so the "Merging" step after the last segment goes away. The output file +is the same either way. + +Segments arrive out of order because the workers run in parallel. unshackle keeps a cursor on +the next segment it needs and appends every segment that is ready in order, then deletes the +segment file. With [`decrypt_segments`](drm.md#decrypt_segments) also on, each segment is +appended after its decryption completes, never before. + +This option is off by default because it disables resume, and because its cost during the +download is not measured. The append reads each segment straight after a worker wrote it, so +the read usually comes from the page cache. On a hard disk, the append to one file while the +workers write many segment files can still make the head seek. If a download gets slower with +the option on, turn it off. + +This applies to DASH video and audio tracks only. Subtitle tracks, HLS tracks and ISM +tracks keep the merge after the download. + +!!! note "No resume" + + A track that merges segment by segment always downloads from the start. + [`continue_downloads`](#continue_downloads) cannot reuse its segments, because unshackle + deletes each segment after it appends it. + +```yaml +merge_segments: true +``` + ## `subtitle` - **Type:** `dict`  ·  **Default:** `{}` diff --git a/docs/reference/configuration/index.md b/docs/reference/configuration/index.md index bf9f7a14..041d0a60 100644 --- a/docs/reference/configuration/index.md +++ b/docs/reference/configuration/index.md @@ -31,7 +31,7 @@ Where the file lives, how unshackle discovers it, and how to edit it: |------|------| | [Directories](directories.md) | `directories`, `filenames`, `services_repo_force` | | [Services & authentication](services.md) | `services`, `credentials`, `firefox_cookies`, `remote_services`, `serve` | -| [Download & processing](download.md) | `dl`, `continue_downloads`, `subtitle`, `audio`, `muxing`, `language_tags`, `dual_multi_mode` | +| [Download & processing](download.md) | `dl`, `continue_downloads`, `merge_segments`, `subtitle`, `audio`, `muxing`, `language_tags`, `dual_multi_mode` | | [Output & naming](output.md) | `output_template`, `tag`, `tag_rules`, `tag_group_name`, `tag_imdb_tmdb`, `chapter_fallback_name`, `unicode_filenames` | | [DRM & CDM](drm.md) | `cdm`, `remote_cdm`, `decryption` | | [Network & proxy](network.md) | `network`, `headers`, `proxy_providers` | diff --git a/tests/core/test_dash_collapse_integration.py b/tests/core/test_dash_collapse_integration.py index c53d0193..133f7cbc 100644 --- a/tests/core/test_dash_collapse_integration.py +++ b/tests/core/test_dash_collapse_integration.py @@ -24,6 +24,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest +from unshackle.core.config import config from unshackle.core.constants import DOWNLOAD_CANCELLED, DownloadCancelled from unshackle.core.manifests import DASH from unshackle.core.tracks.track import DownloadContext @@ -235,3 +236,33 @@ def test_cancelled_sibling_does_not_report_missing_segments(server, tmp_path): # missing segments here would mask the sibling failure that set the flag assert track.path is None assert not ctx.save_path.exists() + + +def test_rolling_merge_appends_in_order_and_leaves_no_segments(server, tmp_path, monkeypatch): + monkeypatch.setattr(config, "merge_segments", True) + server.routes["/init.bin"] = INIT2 + server.routes["/part0.bin"] = PART0 + server.routes["/part1.bin"] = PART1 + track = make_track(mixed_mpd(server), server) + ctx = make_ctx(tmp_path, "rolling.mp4") + + DASH.download_track(track, ctx) + + assert ctx.save_path.read_bytes() == INIT2 + PART0 + PART1 + assert not ctx.save_dir.exists() or not list(ctx.save_dir.glob("*")) + + +def test_rolling_merge_cancel_leaves_no_partial_output(server, tmp_path, monkeypatch): + monkeypatch.setattr(config, "merge_segments", True) + server.routes["/init.bin"] = INIT2 + server.routes["/part0.bin"] = PART0 + server.routes["/part1.bin"] = PART1 + track = make_track(mixed_mpd(server), server) + ctx = make_ctx(tmp_path, "rolling-cancelled.mp4") + + DOWNLOAD_CANCELLED.set() + with pytest.raises(DownloadCancelled): + DASH.download_track(track, ctx) + + assert track.path is None + assert not ctx.save_path.exists() diff --git a/tests/test_rolling_merge.py b/tests/test_rolling_merge.py new file mode 100644 index 00000000..cf089762 --- /dev/null +++ b/tests/test_rolling_merge.py @@ -0,0 +1,51 @@ +import io +import random +from concurrent.futures import Future +from pathlib import Path + +from unshackle.core.manifests.dash import RollingMerge, append_segment + + +def write_segments(save_dir: Path, payloads: list[bytes]) -> int: + save_dir.mkdir() + width = len(str(len(payloads))) + for i, payload in enumerate(payloads): + (save_dir / f"{i:0{width}d}.mp4").write_bytes(payload) + return width + + +def expected(init: bytes, payloads: list[bytes], is_text_subtitle: bool, tmp_path: Path) -> bytes: + scratch = tmp_path / "scratch" + width = write_segments(scratch, payloads) + out = io.BytesIO() + out.write(init) + for i in range(len(payloads)): + append_segment(out, scratch / f"{i:0{width}d}.mp4", is_text_subtitle) + return out.getvalue() + + +def test_out_of_order_segments_match_whole_merge(tmp_path: Path) -> None: + rng = random.Random(7) + payloads = [rng.randbytes(rng.randint(1, 4096)) for _ in range(23)] + init = b"ftyp-moov" + want = expected(init, payloads, False, tmp_path) + + save_dir = tmp_path / "segments" + width = write_segments(save_dir, payloads) + dst = io.BytesIO() + dst.write(init) + merger = RollingMerge(dst, save_dir, width) + + order = list(range(len(payloads))) + rng.shuffle(order) + for index in order: + future: Future = Future() + future.set_result(None) + merger.add(index, future if index % 2 else None) + # the cursor never rests on a segment that is ready: every contiguous run is drained + assert merger.merged not in merger.ready + + assert merger.merged == len(payloads) + assert dst.getvalue() == want + assert dst.getvalue().startswith(init) + assert list(save_dir.iterdir()) == [] diff --git a/unshackle/core/config.py b/unshackle/core/config.py index 0ce9ef50..74e6f828 100644 --- a/unshackle/core/config.py +++ b/unshackle/core/config.py @@ -158,6 +158,7 @@ class Config: self.decryption_map = {} self.decryption = decryption_cfg or "shaka" self.decrypt_segments: bool = bool(kwargs.get("decrypt_segments", False)) + self.merge_segments: bool = bool(kwargs.get("merge_segments", False)) self.theme: str = kwargs.get("theme") or "default" self.set_terminal_bg: bool = kwargs.get("set_terminal_bg", False) diff --git a/unshackle/core/drm/segment_decrypt.py b/unshackle/core/drm/segment_decrypt.py index 575a6bec..93f315db 100644 --- a/unshackle/core/drm/segment_decrypt.py +++ b/unshackle/core/drm/segment_decrypt.py @@ -53,10 +53,13 @@ class SegmentDecrypter: self.pool = ThreadPoolExecutor(max_workers=min(workers or os.cpu_count() or 4, MAX_WORKERS)) self.futures: list[Future] = [] + self._init_bytes: Optional[bytes] = None - def submit(self, segment: Path) -> None: - """Queue one downloaded segment for decryption.""" - self.futures.append(self.pool.submit(self.decrypt_segment, segment)) + def submit(self, segment: Path) -> Future: + """Queue one downloaded segment for decryption and return its future.""" + future = self.pool.submit(self.decrypt_segment, segment) + self.futures.append(future) + return future def decrypt_segment(self, segment: Path) -> None: output = segment.with_suffix(segment.suffix + ".dec") @@ -66,14 +69,24 @@ class SegmentDecrypter: finally: output.unlink(missing_ok=True) + def init_bytes(self) -> bytes: + """Decrypt the init segment on first call and return its bytes. + + Independent of the media segments, so a rolling merge can write it at the head of the + output before any segment lands. + """ + if self._init_bytes is None: + decrypted_init = self.work_dir / "init_decrypted.mp4" + self.run([str(self.init_path), str(decrypted_init)]) + self._init_bytes = decrypted_init.read_bytes() + return self._init_bytes + def finish(self) -> bytes: """Wait for every queued segment, then return the decrypted init segment bytes.""" try: for future in self.futures: future.result() - decrypted_init = self.work_dir / "init_decrypted.mp4" - self.run([str(self.init_path), str(decrypted_init)]) - return decrypted_init.read_bytes() + return self.init_bytes() finally: self.close() diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 6db1fbcf..20ee9b5f 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -7,9 +7,12 @@ import math import re import shutil import sys +from concurrent.futures import Future +from contextlib import ExitStack, closing from copy import deepcopy from functools import lru_cache, partial -from typing import Any, Callable, Optional, Union +from pathlib import Path +from typing import Any, BinaryIO, Callable, Optional, Union from urllib.parse import urljoin, urlparse from uuid import UUID from zlib import crc32 @@ -37,6 +40,61 @@ from unshackle.core.utils.redact import safe_display_url from unshackle.core.utils.xml import load_xml +def append_segment(dst: BinaryIO, segment_file: Path, is_text_subtitle: bool) -> None: + """Append one segment file to the output, with the text subtitle fixups when they apply.""" + if is_text_subtitle: + segment_data = try_ensure_utf8(segment_file.read_bytes()) + segment_data = ( + segment_data.decode("utf8") + .replace("‎", html.unescape("‎")) + .replace("‏", html.unescape("‏")) + .encode("utf8") + ) + dst.write(segment_data) + else: + with open(segment_file, "rb") as src: + shutil.copyfileobj(src, dst, 1024 * 1024) + + +class RollingMerge: + """Appends segments to the output in index order as they arrive, out of order, and unlinks them. + + A cursor holds the next expected index; every `add` records one ready segment and then drains + every contiguous ready segment from the cursor onward. A segment that carries a future (per-segment + decryption) waits for that future before the append, so the output never holds raw bytes. + """ + + def __init__(self, dst: BinaryIO, save_dir: Path, width: int) -> None: + self.dst = dst + self.save_dir = save_dir + self.width = width + self.cursor = 0 + self.ready: dict[int, Optional[Future]] = {} + + def add(self, index: int, future: Optional[Future] = None) -> None: + # ponytail: runs on the downloader's own loop, so a decrypt wait or a catch-up burst + # delays new submissions; move the drain to a writer thread if throughput measurably drops + self.ready[index] = future + while self.cursor in self.ready: + pending = self.ready.pop(self.cursor) + if pending is not None: + pending.result() + segment_file = self.save_dir / f"{self.cursor:0{self.width}d}.mp4" + append_segment(self.dst, segment_file, is_text_subtitle=False) + self.cursor += 1 + # the bytes are already in the output; a handle another process still holds on + # the just-replaced file (Windows) must not fail the track, the segment dir sweep + # after the download removes any leftover + try: + segment_file.unlink() + except OSError: + pass + + @property + def merged(self) -> int: + return self.cursor + + class DASH: def __init__(self, manifest, url: str): if manifest is None: @@ -473,11 +531,14 @@ class DASH: ) use_segment_decrypt = not collapse_single_url and bool(init_data) and can_use(drm, config.decryption) + # subtitle segments are tiny, the post-download merge costs nothing there + use_rolling_merge = not collapse_single_url and config.merge_segments and not isinstance(track, Subtitle) if not collapse_single_url: # per-segment decryption rewrites each segment in place, so a segment kept from an - # earlier run would either be decrypted twice or never at all - can_resume = config.continue_downloads and not use_segment_decrypt + # earlier run would either be decrypted twice or never at all; a rolling merge + # unlinks each segment as it is appended, so there is nothing left to reuse + can_resume = config.continue_downloads and not use_segment_decrypt and not use_rolling_merge # reuse a prior run's segments only when the fingerprint proves the segmentation unchanged digest = resume.fingerprint(track.url, segments) if not (can_resume and resume.reusable(save_dir, digest)): @@ -532,33 +593,72 @@ class DASH: }, ) + is_text_subtitle = ( + not drm and isinstance(track, Subtitle) and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML) + ) + decrypter: Optional[SegmentDecrypter] = None if use_segment_decrypt and init_data: decrypter = SegmentDecrypter(drm, init_data, save_dir.parent, max_workers) + merger: Optional[RollingMerge] = None try: - for status_update in downloader(**downloader_args): - file_downloaded = status_update.get("file_downloaded") - if file_downloaded: - if decrypter: - decrypter.submit(file_downloaded) - events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) - else: - downloaded = status_update.get("downloaded") - if downloaded and downloaded.endswith("/s"): - status_update["downloaded"] = f"DASH {downloaded}" - progress(**status_update) + with ExitStack() as stack: + if use_rolling_merge: + output = stack.enter_context(open(save_path, "wb")) + head = decrypter.init_bytes() if decrypter else init_data + if head: + output.write(head) + merger = RollingMerge(output, save_dir, len(str(len(segments)))) + + # an exception raised in this body (a decrypt failure re-raised by the merge) + # must close the generator now, so its finally aborts the workers + stream = stack.enter_context(closing(downloader(**downloader_args))) + for status_update in stream: + file_downloaded = status_update.get("file_downloaded") + if file_downloaded: + future = decrypter.submit(file_downloaded) if decrypter else None + # listeners get the segment while it still exists; the merge unlinks it + events.emit(events.Types.SEGMENT_DOWNLOADED, track=track, segment=file_downloaded) + if merger: + merger.add(int(file_downloaded.stem), future) + else: + downloaded = status_update.get("downloaded") + if downloaded and downloaded.endswith("/s"): + status_update["downloaded"] = f"DASH {downloaded}" + progress(**status_update) if decrypter: init_data = decrypter.finish() except Exception: if decrypter: decrypter.close() + if merger: + save_path.unlink(missing_ok=True) raise if collapse_single_url: with open(save_path, "r+b") as collapsed: collapsed.truncate(int(segments[-1][1].split("-")[1]) + 1) + elif merger: + for control_file in save_dir.glob("*.!dev"): + control_file.unlink(missing_ok=True) + if merger.merged != len(segments): + error_msg = f"Rolling merge appended {merger.merged} of {len(segments)} segments: {save_dir}" + log_event( + "manifest_dash_rolling_merge_incomplete", + level="ERROR", + message=error_msg, + context={ + "track_id": getattr(track, "id", None), + "track_type": track.__class__.__name__, + "save_dir": str(save_dir), + "merged": merger.merged, + "total_segments": len(segments), + "downloader": "requests", + }, + ) + raise FileNotFoundError(error_msg) else: if not save_dir.exists(): error_msg = f"Output directory does not exist: {save_dir}" @@ -613,29 +713,13 @@ class DASH: ) raise FileNotFoundError(error_msg) - is_text_subtitle = ( - not drm - and isinstance(track, Subtitle) - and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML) - ) with open(save_path, "wb") as f: if init_data: f.write(init_data) if len(segments_to_merge) > 1: progress(downloaded="Merging", completed=0, total=len(segments_to_merge)) for segment_file in segments_to_merge: - if is_text_subtitle: - segment_data = try_ensure_utf8(segment_file.read_bytes()) - segment_data = ( - segment_data.decode("utf8") - .replace("‎", html.unescape("‎")) - .replace("‏", html.unescape("‏")) - .encode("utf8") - ) - f.write(segment_data) - else: - with open(segment_file, "rb") as src: - shutil.copyfileobj(src, f, 1024 * 1024) + append_segment(f, segment_file, is_text_subtitle) segment_file.unlink() progress(advance=1) diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 6aff7909..73e370f6 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -619,7 +619,9 @@ class Track: else: save_dir = save_path.parent - keep_segments = config.continue_downloads and self.descriptor != self.Descriptor.URL + keep_segments = ( + config.continue_downloads and not config.merge_segments and self.descriptor != self.Descriptor.URL + ) def cleanup() -> None: save_path.unlink(missing_ok=True)