From f2273b8e97cf53cf1238de4c65d94c55723080c3 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Tue, 20 Oct 2026 19:53:42 -0600 Subject: [PATCH] feat(dl): keep failed track segments across runs in a resume stash Stash a failed or cancelled track's segments outside the task temp dir and restore them on the next download. Also drop a pre-sized data file with its .!dev control file. --- docs/reference/configuration/download.md | 5 + tests/core/test_resume_across_runs.py | 240 +++++++++++++++++++++++ unshackle/core/temp.py | 85 +++++++- unshackle/core/tracks/track.py | 8 + 4 files changed, 337 insertions(+), 1 deletion(-) create mode 100644 tests/core/test_resume_across_runs.py diff --git a/docs/reference/configuration/download.md b/docs/reference/configuration/download.md index 54922c21..37984972 100644 --- a/docs/reference/configuration/download.md +++ b/docs/reference/configuration/download.md @@ -112,6 +112,11 @@ dl: Keep a failed download's completed segment files so the next download resumes from where it stopped instead of restarting from byte zero. +unshackle keeps the segments of a failed track in a `resume` directory inside the temp +directory, and the next download of the same track takes them back from there. It removes a +kept entry after 7 days, and `unshackle env clear temp` removes them all. A process that is +killed, or a job cancelled over the REST API, keeps nothing. + unshackle reuses only whole, fully downloaded segments. It always discards partially written segments, and no partial-byte (`Range`) resume across runs takes place. Reuse is possible only with proof that the segmentation did not change. The manifest URL, segment count, and byte diff --git a/tests/core/test_resume_across_runs.py b/tests/core/test_resume_across_runs.py new file mode 100644 index 00000000..ff0b8554 --- /dev/null +++ b/tests/core/test_resume_across_runs.py @@ -0,0 +1,240 @@ +"""A failed track's kept segments survive the per-run temp dir and come back in the next run.""" + +from __future__ import annotations + +import os +import time +from functools import partial +from pathlib import Path + +import pytest + +from unshackle.core.config import config +from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, DownloadCancelled +from unshackle.core.temp import ( + RESUME_STASH, + STASH_MAX_AGE, + restore_segments, + resume_stash, + stash_segments, + sweep_task_dirs, + task_temp_dir, +) +from unshackle.core.tracks import Video, resume + +pytestmark = pytest.mark.unit + +DIGEST = "f" * 64 +SEG_NAME = "Video_t1.mp4_segments" + + +@pytest.fixture +def temp_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + root = tmp_path / "temp" + root.mkdir() + monkeypatch.setattr(config.directories, "temp", root) + DOWNLOAD_CANCELLED.clear() + DOWNLOAD_LICENCE_ONLY.clear() + yield root + DOWNLOAD_CANCELLED.clear() + + +def make_track() -> Video: + return Video( + id_="t1", + url="https://example.test/manifest.mpd", + language="en", + codec=Video.Codec.AVC, + range_=Video.Range.SDR, + width=1920, + height=1080, + bitrate=5_000_000, + descriptor=Video.Descriptor.DASH, + ) + + +def run_download(monkeypatch: pytest.MonkeyPatch, parser) -> None: + import requests as requests_lib + + from unshackle.core.manifests import DASH + + monkeypatch.setattr(DASH, "download_track", staticmethod(parser)) + make_track().download( + session=requests_lib.Session(), + prepare_drm=partial(lambda *a, **kw: None), + progress=partial(lambda **kw: None), + ) + + +def seed(ctx): + ctx.save_dir.mkdir(parents=True, exist_ok=True) + (ctx.save_dir / "0.mp4").write_bytes(b"complete segment") + resume.write_sidecar(ctx.save_dir, DIGEST) + + +def failing_parser(track, ctx): + seed(ctx) + (ctx.save_dir / "1.mp4.p.!dev").write_bytes(b"partial") + raise RuntimeError("simulated crash") + + +def test_failed_track_resumes_in_next_run(temp_root, monkeypatch): + monkeypatch.setattr(config, "continue_downloads", True) + + with task_temp_dir(): + with pytest.raises(RuntimeError): + run_download(monkeypatch, failing_parser) + + stashed = temp_root / RESUME_STASH / SEG_NAME + assert (stashed / "0.mp4").read_bytes() == b"complete segment" + assert resume.sidecar_path(stashed).exists() + + seen = {} + + def second_parser(track, ctx): + seen["bytes"] = (ctx.save_dir / "0.mp4").read_bytes() + seen["reusable"] = resume.reusable(ctx.save_dir, DIGEST) + ctx.save_path.write_bytes(b"output") + track.path = ctx.save_path + + with task_temp_dir(): + run_download(monkeypatch, second_parser) + + assert seen == {"bytes": b"complete segment", "reusable": True} + assert not stashed.exists() + + +def test_flag_off_stashes_nothing(temp_root, monkeypatch): + monkeypatch.setattr(config, "continue_downloads", False) + + with task_temp_dir(): + with pytest.raises(RuntimeError): + run_download(monkeypatch, failing_parser) + + assert not (temp_root / RESUME_STASH).exists() + + +def test_partials_are_not_stashed(temp_root, monkeypatch): + monkeypatch.setattr(config, "continue_downloads", True) + + with task_temp_dir(): + with pytest.raises(RuntimeError): + run_download(monkeypatch, failing_parser) + + assert (temp_root / RESUME_STASH / SEG_NAME / "0.mp4").exists() + assert not list((temp_root / RESUME_STASH).rglob("*.!dev")) + + +def test_existing_stash_target_is_left_alone(temp_root): + with task_temp_dir() as task_dir: + target = temp_root / RESUME_STASH / SEG_NAME + target.mkdir(parents=True) + (target / "old.mp4").write_bytes(b"old") + save_dir = task_dir / SEG_NAME + save_dir.mkdir() + (save_dir / "new.mp4").write_bytes(b"new") + + stash_segments(save_dir) + + assert (target / "old.mp4").read_bytes() == b"old" + assert not (target / "new.mp4").exists() + assert save_dir.exists() + + +def test_restore_is_a_noop_without_stash_or_with_present_dir(temp_root): + with task_temp_dir() as task_dir: + save_dir = task_dir / SEG_NAME + restore_segments(save_dir) + assert not save_dir.exists() + + target = temp_root / RESUME_STASH / SEG_NAME + target.mkdir(parents=True) + (target / "a.mp4").write_bytes(b"a") + save_dir.mkdir() + + restore_segments(save_dir) + + assert (target / "a.mp4").exists() + assert not (save_dir / "a.mp4").exists() + + +def test_rename_failure_never_replaces_download_error(temp_root, monkeypatch): + monkeypatch.setattr(config, "continue_downloads", True) + + def deny(self, target): + raise PermissionError("locked") + + with task_temp_dir(): + monkeypatch.setattr(Path, "rename", deny) + with pytest.raises(RuntimeError, match="simulated crash"): + run_download(monkeypatch, failing_parser) + + +def test_sweep_removes_only_old_stash_entries(temp_root): + stash = temp_root / RESUME_STASH + old, fresh = stash / "old_segments", stash / "fresh_segments" + old.mkdir(parents=True) + fresh.mkdir() + old_sidecar = stash / "old_segments.resume.json" + old_sidecar.write_text("{}") + stamp = time.time() - STASH_MAX_AGE - 60 + for path in (old, old_sidecar): + os.utime(path, (stamp, stamp)) + + sweep_task_dirs(temp_root) + + assert not old.exists() + assert not old_sidecar.exists() + assert fresh.exists() + + +def test_no_stash_outside_task_dir(temp_root): + assert resume_stash() is None + + +def test_cancelled_track_is_stashed(temp_root, monkeypatch): + monkeypatch.setattr(config, "continue_downloads", True) + + def cancelled_parser(track, ctx): + seed(ctx) + raise DownloadCancelled + + with task_temp_dir(): + run_download(monkeypatch, cancelled_parser) + + stashed = temp_root / RESUME_STASH / SEG_NAME + assert (stashed / "0.mp4").read_bytes() == b"complete segment" + assert resume.sidecar_path(stashed).exists() + + +def test_presized_data_file_is_dropped_with_its_control_file(temp_root, monkeypatch): + monkeypatch.setattr(config, "continue_downloads", True) + + def parser(track, ctx): + seed(ctx) + (ctx.save_dir / "1.mp4").write_bytes(b"\0" * 8) + (ctx.save_dir / "1.mp4.!dev").write_bytes(b"") + raise RuntimeError("simulated crash") + + with task_temp_dir(): + with pytest.raises(RuntimeError): + run_download(monkeypatch, parser) + + stashed = temp_root / RESUME_STASH / SEG_NAME + assert (stashed / "0.mp4").exists() + assert not (stashed / "1.mp4").exists() + assert not (stashed / "1.mp4.!dev").exists() + + +def test_stash_refreshes_mtime(temp_root): + with task_temp_dir() as task_dir: + save_dir = task_dir / SEG_NAME + save_dir.mkdir() + resume.write_sidecar(save_dir, DIGEST) + stamp = time.time() - STASH_MAX_AGE - 60 + os.utime(save_dir, (stamp, stamp)) + + stash_segments(save_dir) + + target = temp_root / RESUME_STASH / SEG_NAME + assert time.time() - target.stat().st_mtime < 60 diff --git a/unshackle/core/temp.py b/unshackle/core/temp.py index d6d45a99..655329a9 100644 --- a/unshackle/core/temp.py +++ b/unshackle/core/temp.py @@ -1,6 +1,9 @@ from __future__ import annotations +import contextlib import functools +import logging +import os import shutil import time import uuid @@ -15,6 +18,10 @@ from unshackle.core.config import config TASK_PREFIX = "task_" LOCK_NAME = ".lock" STALE_GRACE = 60.0 +RESUME_STASH = "resume" +STASH_MAX_AGE = 7 * 24 * 3600 + +log = logging.getLogger("temp") P = ParamSpec("P") R = TypeVar("R") @@ -40,13 +47,89 @@ def is_stale(task_dir: Path) -> bool: def sweep_task_dirs(root: Path) -> None: - """Remove task dirs left by processes that died without cleaning up.""" + """Remove task dirs left by processes that died without cleaning up, and resume stash entries past their age.""" if not root.is_dir(): return for entry in root.iterdir(): if entry.is_dir() and not entry.is_symlink() and entry.name.startswith(TASK_PREFIX) and is_stale(entry): shutil.rmtree(entry, ignore_errors=True) + stash = root / RESUME_STASH + if stash.is_dir(): + cutoff = time.time() - STASH_MAX_AGE + for entry in stash.iterdir(): + try: + if entry.stat().st_mtime > cutoff: + continue + except OSError: + continue + if entry.is_dir() and not entry.is_symlink(): + shutil.rmtree(entry, ignore_errors=True) + else: + entry.unlink(missing_ok=True) + + +def resume_stash() -> Optional[Path]: + """The dir beside the task dirs that keeps a failed track's segments for the next run. + + None outside a task dir: there config.directories.temp itself outlives the run. + """ + temp = config.directories.temp + if not temp.name.startswith(TASK_PREFIX): + return None + return temp.parent / RESUME_STASH + + +def _move_with_sidecar(source: Path, target: Path) -> bool: + """Rename a segment dir, then its resume sidecar. False if the dir rename fails; a sidecar failure only logs.""" + from unshackle.core.tracks import resume + + try: + source.rename(target) + except OSError as e: + log.debug("Could not move %s: %s", source, e) + return False + sidecar = resume.sidecar_path(source) + try: + if sidecar.exists(): + sidecar.replace(resume.sidecar_path(target)) + except OSError as e: + log.debug("Could not move %s: %s", sidecar, e) + return True + + +def stash_segments(save_dir: Path) -> None: + """Move save_dir and its resume sidecar into the stash. Leave them in place on any OS error.""" + from unshackle.core.tracks import resume + + stash = resume_stash() + if stash is None or not save_dir.is_dir(): + return + target = stash / save_dir.name + try: + stash.mkdir(parents=True, exist_ok=True) + if target.exists(): + return + except OSError as e: + log.debug("Could not stash %s: %s", save_dir, e) + return + if _move_with_sidecar(save_dir, target): + # rename keeps the old mtime, and the sweep ages entries by it + for path in (target, resume.sidecar_path(target)): + with contextlib.suppress(OSError): + os.utime(path) + + +def restore_segments(save_dir: Path) -> None: + """Move a stashed save_dir and its sidecar back to save_dir, if the stash has them and save_dir is free.""" + stash = resume_stash() + if stash is None: + return + source = stash / save_dir.name + if source.is_symlink() or not source.is_dir() or save_dir.exists(): + return + _move_with_sidecar(source, save_dir) + @contextmanager def task_temp_dir(task_id: Optional[str] = None) -> Iterator[Path]: diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index acb8a76f..e2dc0e62 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -32,6 +32,7 @@ from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine from unshackle.core.drm.verify import decrypt_track from unshackle.core.events import events from unshackle.core.session import MAX_BACKOFF, RnetSession, utf8_text +from unshackle.core.temp import restore_segments, stash_segments from unshackle.core.tracks import resume from unshackle.core.utilities import get_boxes, log_event, try_ensure_utf8 from unshackle.core.utils import run_proxy @@ -760,6 +761,7 @@ class Track: if keep_segments: if save_dir.exists(): for partial in save_dir.rglob("*.!dev"): + partial.with_name(partial.name.removesuffix(".!dev")).unlink(missing_ok=True) partial.unlink(missing_ok=True) else: if save_dir.exists(): @@ -774,6 +776,8 @@ class Track: # completed segments are reusable once the parser proves the segmentation # unchanged (resume sidecar); partial .!dev files never survive a run boundary + if keep_segments: + restore_segments(save_dir) cleanup() try: @@ -925,6 +929,8 @@ class Track: except DownloadCancelled: try: cleanup() + if keep_segments: + stash_segments(save_dir) except OSError: pass progress(downloaded="[yellow]SKIPPED") @@ -932,6 +938,8 @@ class Track: except (Exception, KeyboardInterrupt): if not DOWNLOAD_LICENCE_ONLY.is_set(): cleanup() + if keep_segments: + stash_segments(save_dir) raise if DOWNLOAD_CANCELLED.is_set():