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

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

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

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]:

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():