diff --git a/tests/core/test_task_temp.py b/tests/core/test_task_temp.py new file mode 100644 index 00000000..69aceaa0 --- /dev/null +++ b/tests/core/test_task_temp.py @@ -0,0 +1,174 @@ +"""Tests for per-task temp directories. + +Each run gets a private dir that is removed on every exit path. The sweep reclaims +dirs from hard-killed processes via lock liveness, and `env clear temp` skips the +dir of a task that is still running. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import pytest +from filelock import FileLock + +from unshackle.commands.env import clear_directory +from unshackle.core.config import config +from unshackle.core.temp import LOCK_NAME, TASK_PREFIX, is_stale, sweep_task_dirs, task_temp_dir + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def temp_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + root = tmp_path / "temp" + root.mkdir() + monkeypatch.setattr(config.directories, "temp", root) + return root + + +def backdate(path: Path) -> None: + stamp = time.time() - 120 + os.utime(path, (stamp, stamp)) + + +def make_task_dir(root: Path, name: str = "task_deadbeef", lock: bool = True) -> Path: + task_dir = root / name + task_dir.mkdir() + if lock: + (task_dir / LOCK_NAME).touch() + return task_dir + + +def test_task_dir_created_repointed_and_restored(temp_root: Path) -> None: + with task_temp_dir() as task_dir: + assert task_dir.parent == temp_root + assert task_dir.name.startswith(TASK_PREFIX) + assert config.directories.temp == task_dir + assert config.directories.temp == temp_root + assert not task_dir.exists() + + +def test_task_dir_removed_on_exception(temp_root: Path) -> None: + with pytest.raises(ValueError): + with task_temp_dir() as task_dir: + (task_dir / "partial.mp4").write_bytes(b"x") + raise ValueError("boom") + assert not task_dir.exists() + assert config.directories.temp == temp_root + + +def test_task_dir_removed_on_keyboard_interrupt(temp_root: Path) -> None: + with pytest.raises(KeyboardInterrupt): + with task_temp_dir() as task_dir: + raise KeyboardInterrupt + assert not task_dir.exists() + assert config.directories.temp == temp_root + + +def test_sweep_removes_unlocked_stale_dir(temp_root: Path) -> None: + task_dir = make_task_dir(temp_root) + (task_dir / "leftover.mp4").write_bytes(b"x") + backdate(task_dir) + sweep_task_dirs(temp_root) + assert not task_dir.exists() + + +def test_sweep_skips_locked_dir(temp_root: Path) -> None: + task_dir = make_task_dir(temp_root) + lock = FileLock(task_dir / LOCK_NAME) + lock.acquire() + backdate(task_dir) + try: + assert is_stale(task_dir) is False + sweep_task_dirs(temp_root) + assert task_dir.exists() + finally: + lock.release() + + +def test_sweep_skips_young_unlocked_dir(temp_root: Path) -> None: + task_dir = make_task_dir(temp_root) + assert is_stale(task_dir) is False + sweep_task_dirs(temp_root) + assert task_dir.exists() + + +def test_sweep_removes_stale_dir_without_lock_file(temp_root: Path) -> None: + task_dir = make_task_dir(temp_root, lock=False) + backdate(task_dir) + assert is_stale(task_dir) is True + sweep_task_dirs(temp_root) + assert not task_dir.exists() + + +def test_sweep_ignores_non_task_entries(temp_root: Path) -> None: + loose_file = temp_root / "leftover.mp4" + loose_file.write_bytes(b"x") + other_dir = temp_root / "not_a_task" + other_dir.mkdir() + backdate(loose_file) + backdate(other_dir) + sweep_task_dirs(temp_root) + assert loose_file.exists() + assert other_dir.exists() + + +@pytest.mark.skipif(os.name == "nt", reason="symlink creation needs privileges on Windows") +def test_sweep_never_follows_symlinked_task_dir(temp_root: Path, tmp_path: Path) -> None: + target = tmp_path / "victim" + target.mkdir() + (target / "precious.mkv").write_bytes(b"data") + link = temp_root / "task_evil" + link.symlink_to(target) + backdate(target) + sweep_task_dirs(temp_root) + assert (target / "precious.mkv").exists() + assert not (target / LOCK_NAME).exists() + + +@pytest.mark.skipif(os.name == "nt", reason="symlink creation needs privileges on Windows") +def test_clear_directory_unlinks_symlinked_dir_without_following(temp_root: Path, tmp_path: Path) -> None: + target = tmp_path / "victim" + target.mkdir() + (target / "precious.mkv").write_bytes(b"data") + link = temp_root / "task_evil" + link.symlink_to(target) + backdate(target) + files_count, _ = clear_directory(temp_root) + assert not link.exists() + assert (target / "precious.mkv").exists() + assert files_count == 1 + + +def test_clear_directory_skips_locked_task_dir(temp_root: Path) -> None: + task_dir = make_task_dir(temp_root) + (task_dir / "in_flight.mp4").write_bytes(b"xyz") + lock = FileLock(task_dir / LOCK_NAME) + lock.acquire() + backdate(task_dir) + loose_file = temp_root / "leftover.mp4" + loose_file.write_bytes(b"1234") + try: + files_count, freed_bytes = clear_directory(temp_root) + finally: + lock.release() + assert task_dir.exists() + assert (task_dir / "in_flight.mp4").exists() + assert not loose_file.exists() + assert files_count == 1 + assert freed_bytes == 4 + + +def test_clear_directory_removes_stale_task_dir(temp_root: Path) -> None: + task_dir = make_task_dir(temp_root) + (task_dir / "orphan.mp4").write_bytes(b"12345") + backdate(task_dir) + files_count, freed_bytes = clear_directory(temp_root) + assert not task_dir.exists() + assert temp_root.is_dir() + # glob("**/*") never matches dotfiles, so .lock is not counted + assert files_count == 1 + assert freed_bytes == 5 diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index a8615afb..0902518b 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -63,6 +63,7 @@ from unshackle.core.providers.tvdb import SEASON_TYPES from unshackle.core.proxies import Basic, ExpressVPN, Gluetun, Hola, NordVPN, ProtonVPN, SurfsharkVPN, WindscribeVPN from unshackle.core.service import Service from unshackle.core.services import Services +from unshackle.core.temp import with_task_temp from unshackle.core.title_cacher import get_account_hash from unshackle.core.titles import Movie, Movies, Music, Series, Song, Title_T from unshackle.core.titles.episode import Episode @@ -1246,6 +1247,7 @@ class dl: # able to keep `self` as the first positional self.cli._result_callback = self.result + @with_task_temp def result( self, service: Service, diff --git a/unshackle/commands/env.py b/unshackle/commands/env.py index 9e193ac9..f4060677 100644 --- a/unshackle/commands/env.py +++ b/unshackle/commands/env.py @@ -15,6 +15,7 @@ from unshackle.core.config import POSSIBLE_CONFIG_PATHS, config, config_path from unshackle.core.console import console from unshackle.core.constants import context_settings from unshackle.core.services import Services +from unshackle.core.temp import TASK_PREFIX, is_stale def get_dependencies() -> list[dict]: @@ -94,18 +95,36 @@ def get_dependencies() -> list[dict]: def clear_directory(path: Path) -> tuple[int, int]: - """Delete a directory's contents, returning (files_removed, freed_bytes); recreates the dir.""" + """Delete a directory's contents, returning (files_removed, freed_bytes); recreates the dir. + + Skips task directories that belong to a running download. + """ files_count = 0 freed_bytes = 0 if path.exists(): - for entry in path.glob("**/*"): - if entry.is_file(): + for entry in path.iterdir(): + is_real_dir = entry.is_dir() and not entry.is_symlink() + if is_real_dir and entry.name.startswith(TASK_PREFIX) and not is_stale(entry): + continue + if is_real_dir: + for child in entry.glob("**/*"): + if child.is_file(): + files_count += 1 + try: + freed_bytes += child.stat().st_size + except OSError: + pass + shutil.rmtree(entry, ignore_errors=True) + else: files_count += 1 try: - freed_bytes += entry.stat().st_size + freed_bytes += entry.lstat().st_size + except OSError: + pass + try: + entry.unlink(missing_ok=True) except OSError: pass - shutil.rmtree(path, ignore_errors=True) path.mkdir(parents=True, exist_ok=True) return files_count, freed_bytes diff --git a/unshackle/core/temp.py b/unshackle/core/temp.py new file mode 100644 index 00000000..b70d2573 --- /dev/null +++ b/unshackle/core/temp.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import functools +import shutil +import time +import uuid +from contextlib import contextmanager +from pathlib import Path +from typing import Callable, Iterator, Optional, ParamSpec, TypeVar + +from filelock import FileLock, Timeout + +from unshackle.core.config import config + +TASK_PREFIX = "task_" +LOCK_NAME = ".lock" +STALE_GRACE = 60.0 + +P = ParamSpec("P") +R = TypeVar("R") + + +def is_stale(task_dir: Path) -> bool: + """A task dir is stale when nothing holds its lock and it is past the setup grace.""" + try: + if time.time() - task_dir.stat().st_mtime < STALE_GRACE: + return False + except OSError: + return False + lock_path = task_dir / LOCK_NAME + if not lock_path.exists(): + return True + lock = FileLock(lock_path, timeout=0) + try: + lock.acquire() + except (Timeout, OSError): + return False + lock.release() + return True + + +def sweep_task_dirs(root: Path) -> None: + """Remove task dirs left by processes that died without cleaning up.""" + 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) + + +@contextmanager +def task_temp_dir(task_id: Optional[str] = None) -> Iterator[Path]: + """Point config.directories.temp at a private dir for this run; remove it on any exit.""" + root = config.directories.temp + root.mkdir(parents=True, exist_ok=True) + sweep_task_dirs(root) + + task_dir = root / f"{TASK_PREFIX}{task_id or uuid.uuid4().hex[:12]}" + task_dir.mkdir(parents=True, exist_ok=True) + lock = FileLock(task_dir / LOCK_NAME) + lock.acquire() + config.directories.temp = task_dir + try: + yield task_dir + finally: + config.directories.temp = root + lock.release() + shutil.rmtree(task_dir, ignore_errors=True) + + +def with_task_temp(fn: Callable[P, R]) -> Callable[P, R]: + """Run the wrapped callable inside its own task temp dir.""" + + @functools.wraps(fn) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + with task_temp_dir(): + return fn(*args, **kwargs) + + return wrapper