From 2f28aa8e643ae5dde1448196fd116ae1aaf094c6 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sat, 5 Sep 2026 12:52:02 -0600 Subject: [PATCH] fix(dl): end a cancelled track by raising instead of returning short --- tests/core/test_dash_collapse_integration.py | 24 +++++++++++++-- tests/downloaders/test_cancel_hardening.py | 32 +++++++++++++++++++- tests/downloaders/test_multiprocess.py | 13 +++++--- tests/test_downloader_leaks.py | 11 ++++--- unshackle/core/constants.py | 10 ++++++ unshackle/core/downloaders/requests.py | 10 +++++- unshackle/core/tracks/track.py | 11 ++++++- 7 files changed, 96 insertions(+), 15 deletions(-) diff --git a/tests/core/test_dash_collapse_integration.py b/tests/core/test_dash_collapse_integration.py index bbf8cf4e..c53d0193 100644 --- a/tests/core/test_dash_collapse_integration.py +++ b/tests/core/test_dash_collapse_integration.py @@ -9,7 +9,10 @@ localhost server, asserting that: - a SegmentList whose ranges collapse fetches the parent resource whole in one direct (Range-less) request and writes a byte-correct file at ``save_path``, leaving no segment dir or ``.!dev`` artifacts, and -- a track with mixed segment URLs still takes the per-segment download-and-merge path. +- a track with mixed segment URLs still takes the per-segment download-and-merge path, and +- that same segmented path raises ``DownloadCancelled`` when a sibling track's failure has + already set ``DOWNLOAD_CANCELLED``, instead of treating the segments the cancel left + behind as a fault of its own. """ from __future__ import annotations @@ -21,7 +24,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest -from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.constants import DOWNLOAD_CANCELLED, DownloadCancelled from unshackle.core.manifests import DASH from unshackle.core.tracks.track import DownloadContext @@ -215,3 +218,20 @@ def test_non_collapse_takes_segmented_merge_path(server, tmp_path): assert {"/part0.bin", "/part1.bin"} <= fetched # segment files are consumed and the segment dir is cleaned up after merge assert not list(ctx.save_dir.glob("*")) if ctx.save_dir.exists() else True + + +def test_cancelled_sibling_does_not_report_missing_segments(server, tmp_path): + 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, "cancelled.mp4") + + DOWNLOAD_CANCELLED.set() + with pytest.raises(DownloadCancelled): + DASH.download_track(track, ctx) + + # the workers returned early, so there is nothing to merge; a FileNotFoundError about the + # missing segments here would mask the sibling failure that set the flag + assert track.path is None + assert not ctx.save_path.exists() diff --git a/tests/downloaders/test_cancel_hardening.py b/tests/downloaders/test_cancel_hardening.py index 95b3d215..69ac6f59 100644 --- a/tests/downloaders/test_cancel_hardening.py +++ b/tests/downloaders/test_cancel_hardening.py @@ -1,6 +1,6 @@ """Cancel/abort hardening for the downloader. -Pins two teardown-safety invariants: +Pins three teardown-safety invariants: - ``dispatch_parts`` must NOT finalize when the process-global ``DOWNLOAD_CANCELLED`` fires mid-download: part workers return silently keeping partials, so every future @@ -9,6 +9,10 @@ Pins two teardown-safety invariants: - A worker parked in the retry backoff wait (which can reach MAX_BACKOFF) must wake and exit as soon as the batch abort is set, instead of sleeping out the full delay and stalling shutdown or the merge. +- A segment batch that starts with a cross-track cancel set, or meets one mid-flight, must + end by raising ``DownloadCancelled`` in adaptive mode as well as fixed. Adaptive's tail + window holds back boost candidates, and the cancel keeps ``seg_done`` short, so without + the raise it busy-loops forever on a batch it will never finish. """ import importlib @@ -19,6 +23,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest from requests import Session +from unshackle.core.constants import DownloadCancelled + dl = importlib.import_module("unshackle.core.downloaders.requests") pytestmark = pytest.mark.unit @@ -143,3 +149,27 @@ def test_backoff_wait_exits_promptly_on_batch_abort(server, tmp_path, monkeypatc assert not worker.is_alive() # generous bound: far below the 30s nap, so the wait was interrupted, not slept out assert time.monotonic() - start < 10 + + +@pytest.mark.parametrize("adaptive", [False, True]) +def test_cancelled_batch_raises_instead_of_finishing_or_hanging(server, tmp_path, adaptive): + dl.DOWNLOAD_CANCELLED.set() + urls = [f"{url(server)}?i={i}" for i in range(40)] + result: list = [] + + def run() -> None: + try: + for _ in dl.requests(urls, tmp_path, "{i}.mp4", max_workers=16, adaptive=adaptive): + pass + result.append(None) + except BaseException as exc: # noqa: BLE001 (the type is the assertion) + result.append(exc) + + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(timeout=15) + + assert not worker.is_alive(), "requests() never returned on a cancelled batch" + assert isinstance(result[0], DownloadCancelled) + # returning short here would hand the caller an empty segment dir to merge + assert not list(tmp_path.iterdir()) diff --git a/tests/downloaders/test_multiprocess.py b/tests/downloaders/test_multiprocess.py index 22b781ac..d5688063 100644 --- a/tests/downloaders/test_multiprocess.py +++ b/tests/downloaders/test_multiprocess.py @@ -26,6 +26,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest +from unshackle.core.constants import DownloadCancelled + dl = importlib.import_module("unshackle.core.downloaders.requests") SEG_REPEAT = 32 # 8-digit index * 32 = 256-byte body, distinct per index @@ -139,11 +141,12 @@ def test_multiprocess_cancel_terminates_children(server, tmp_path): baseline = len(multiprocessing.active_children()) start = time.time() progressed = 0 - for event in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2): - if event.get("advance") or event.get("written"): - progressed += 1 - if progressed >= 2: - dl.DOWNLOAD_CANCELLED.set() # sibling-track cancel: parent must tear the children down + with pytest.raises(DownloadCancelled): + for event in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2): + if event.get("advance") or event.get("written"): + progressed += 1 + if progressed >= 2: + dl.DOWNLOAD_CANCELLED.set() # sibling-track cancel: parent must tear the children down elapsed = time.time() - start # cancel stopped the work: nowhere near all segments finished, and no hang diff --git a/tests/test_downloader_leaks.py b/tests/test_downloader_leaks.py index 49580a69..0a2cc267 100644 --- a/tests/test_downloader_leaks.py +++ b/tests/test_downloader_leaks.py @@ -19,7 +19,7 @@ from typing import Any, Iterator import pytest -from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.constants import DOWNLOAD_CANCELLED, DownloadCancelled from unshackle.core.downloaders.requests import requests as download_batch dl_mod = importlib.import_module("unshackle.core.downloaders.requests") @@ -122,14 +122,15 @@ def test_failed_batch_frees_handles_even_after_flag_clear(server: _Server, tmp_p def test_external_cancel_frees_handles(server: _Server, tmp_path: Path) -> None: """DOWNLOAD_CANCELLED set mid-stream (sibling-track failure / SIGINT path): the batch - must wind down without leaving locked handles or live workers.""" + must raise DownloadCancelled and wind down without leaving locked handles or live workers.""" baseline = pool_threads() host, port = str(server.server_address[0]), int(server.server_address[1]) urls = [f"http://{host}:{port}/slow/{i}" for i in range(4)] - for event in download_batch(urls=urls, output_dir=tmp_path, filename="seg_{i}.bin", max_workers=4): - if "advance" in event or "written" in event: - DOWNLOAD_CANCELLED.set() + with pytest.raises(DownloadCancelled): + for event in download_batch(urls=urls, output_dir=tmp_path, filename="seg_{i}.bin", max_workers=4): + if "advance" in event or "written" in event: + DOWNLOAD_CANCELLED.set() DOWNLOAD_CANCELLED.clear() assert_drained_and_unlocked(tmp_path, baseline) diff --git a/unshackle/core/constants.py b/unshackle/core/constants.py index 0c3a08fe..51ad0f3b 100644 --- a/unshackle/core/constants.py +++ b/unshackle/core/constants.py @@ -4,6 +4,16 @@ from typing import TypeVar, Union DOWNLOAD_CANCELLED = Event() DOWNLOAD_LICENCE_ONLY = Event() + +class DownloadCancelled(Exception): + """A track stopped because ``DOWNLOAD_CANCELLED`` was set, not because it itself failed. + + The downloader raises this instead of returning short. A short return sends the track on + to its merge or decrypt step, which then fails on the segments the cancel left behind. + That fault hides the sibling failure that is the real error. + """ + + DRM_SORT_MAP = ["ClearKey", "Widevine"] LANGUAGE_MAX_DISTANCE = 5 # this is max to be considered "same", e.g., en, en-US, en-AU LANGUAGE_EXACT_DISTANCE = 0 # exact match only, no variants diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 8288f4c7..97f5be9f 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -25,7 +25,7 @@ from requests import Session from requests.adapters import HTTPAdapter, Retry from rich import filesize -from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.constants import DOWNLOAD_CANCELLED, DownloadCancelled from unshackle.core.utilities import get_debug_logger, get_extension MAX_ATTEMPTS = 5 @@ -1507,6 +1507,8 @@ def requests( processes=processes, debug_logger=debug_logger, ) + if DOWNLOAD_CANCELLED.is_set(): + raise DownloadCancelled() return if debug_logger: debug_logger.log( @@ -1903,6 +1905,9 @@ def requests( yield dict(advance=pending_advance) pending_advance = 0 + if DOWNLOAD_CANCELLED.is_set(): + remaining.clear() + now = time.time() target = max_workers @@ -2032,6 +2037,9 @@ def requests( download_speed = math.ceil(total_bytes / elapsed) yield dict(downloaded=f"{filesize.decimal(download_speed)}/s") + if DOWNLOAD_CANCELLED.is_set(): + raise DownloadCancelled() + if debug_logger: debug_logger.log( level="DEBUG", diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index ff7a9d01..06c6a513 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -21,7 +21,7 @@ from requests.adapters import HTTPAdapter, Retry from unshackle.core import binaries from unshackle.core.cdm.detect import is_playready_cdm, is_widevine_cdm from unshackle.core.config import config -from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY +from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, DownloadCancelled from unshackle.core.downloaders import requests from unshackle.core.drm import DRM_T, ClearKeyCENC, PlayReady, Widevine from unshackle.core.events import events @@ -743,10 +743,19 @@ class Track: DOWNLOAD_CANCELLED.set() progress(downloaded="[yellow]CANCELLED") raise + except DownloadCancelled: + raise except Exception: DOWNLOAD_CANCELLED.set() progress(downloaded="[red]FAILED") raise + except DownloadCancelled: + try: + cleanup() + except OSError: + pass + progress(downloaded="[yellow]SKIPPED") + return except (Exception, KeyboardInterrupt): if not DOWNLOAD_LICENCE_ONLY.is_set(): cleanup()