484 lines
17 KiB
Python
484 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
from imagefind.config import Settings
|
|
from imagefind.remote import RcloneManager
|
|
from imagefind.remote_cache import RemoteMediaCache
|
|
from imagefind.resources import ResourceGovernor
|
|
from imagefind.usage import StorageUsageService
|
|
|
|
|
|
class _SettingsDb:
|
|
def __init__(self):
|
|
self.values = {}
|
|
|
|
def setting(self, key, default=None):
|
|
return self.values.get(key, default)
|
|
|
|
def set_setting(self, key, value):
|
|
self.values[key] = value
|
|
|
|
|
|
class _RemoteSource:
|
|
def remote_access(self, _source_id: str, _key: str):
|
|
return "https://media.test/video.mp4", "user", "password", True
|
|
|
|
|
|
def _sample(some: float, full: float) -> dict:
|
|
return {
|
|
"cpu_percent": 0.0,
|
|
"memory_total_bytes": 8 * 1024**3,
|
|
"memory_available_bytes": 6 * 1024**3,
|
|
"disk_available_bytes": 100 * 1024**3,
|
|
"io": {
|
|
"supported": True,
|
|
"some_avg10": some,
|
|
"some_avg60": some,
|
|
"some_avg300": some,
|
|
"full_avg10": full,
|
|
"full_avg60": full,
|
|
"full_avg300": full,
|
|
},
|
|
}
|
|
|
|
|
|
def test_io_pressure_profiles_use_hysteresis(tmp_path: Path):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
db = _SettingsDb()
|
|
governor = ResourceGovernor(db, settings)
|
|
governor.sample = lambda: _sample(31, 11) # type: ignore[method-assign]
|
|
|
|
for generation in range(1, 4):
|
|
governor._sample_generation = generation
|
|
state = governor._update_io_state(governor.sample())
|
|
assert state == "paused"
|
|
assert "磁盘 I/O" in governor.pressure_reason(lane="ai", running=True)
|
|
|
|
governor.sample = lambda: _sample(2, 0) # type: ignore[method-assign]
|
|
for generation in range(4, 9):
|
|
governor._sample_generation = generation
|
|
state = governor._update_io_state(governor.sample())
|
|
assert state == "normal"
|
|
|
|
db.set_setting("resource_profile", "quiet")
|
|
assert governor.profile() == "quiet"
|
|
|
|
|
|
def test_io_pressure_uses_data_volume_instead_of_system_psi(tmp_path: Path):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
governor = ResourceGovernor(_SettingsDb(), settings)
|
|
sample = _sample(95, 80)
|
|
sample["io"].update(
|
|
{
|
|
"scope": "data_volume",
|
|
"volume_supported": True,
|
|
"volume_sampled": True,
|
|
"volume_utilization_percent": 4.0,
|
|
"volume_in_flight": 0,
|
|
}
|
|
)
|
|
|
|
for generation in range(1, 5):
|
|
governor._sample_generation = generation
|
|
state = governor._update_io_state(sample)
|
|
assert state == "normal"
|
|
|
|
sample["io"]["volume_utilization_percent"] = 92.0
|
|
for generation in range(5, 8):
|
|
governor._sample_generation = generation
|
|
state = governor._update_io_state(sample)
|
|
assert state == "paused"
|
|
|
|
|
|
def test_data_volume_block_device_is_resolved_from_mountinfo(tmp_path: Path):
|
|
data = tmp_path / "volume" / "app" / "data"
|
|
data.mkdir(parents=True)
|
|
mountinfo = tmp_path / "mountinfo"
|
|
mountinfo.write_text(
|
|
f"10 1 8:1 / {tmp_path / 'volume'} rw - ext4 /dev/test rw\n",
|
|
encoding="utf-8",
|
|
)
|
|
sysfs = tmp_path / "sys" / "dev" / "block"
|
|
stat = sysfs / "8:1" / "stat"
|
|
stat.parent.mkdir(parents=True)
|
|
stat.write_text("1 0 0 0 2 0 0 0 3 456 0", encoding="utf-8")
|
|
|
|
result = ResourceGovernor._data_volume_io_values(
|
|
data,
|
|
mountinfo_path=mountinfo,
|
|
sysfs_root=sysfs,
|
|
)
|
|
|
|
assert result == {
|
|
"device": "8:1",
|
|
"mount_point": str(tmp_path / "volume"),
|
|
"in_flight": 3,
|
|
"io_ticks_ms": 456,
|
|
}
|
|
|
|
|
|
def test_remote_media_is_downloaded_once_and_reused(tmp_path: Path, monkeypatch):
|
|
payload = b"0123456789" * 1024
|
|
requests = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
range_header = request.headers.get("range")
|
|
if range_header:
|
|
offset = int(range_header.removeprefix("bytes=").removesuffix("-"))
|
|
return httpx.Response(
|
|
206,
|
|
headers={"Content-Range": f"bytes {offset}-{len(payload) - 1}/{len(payload)}"},
|
|
content=payload[offset:],
|
|
)
|
|
return httpx.Response(200, content=payload)
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
client_type = httpx.Client
|
|
monkeypatch.setattr(
|
|
"imagefind.remote_cache.httpx.Client",
|
|
lambda **kwargs: client_type(transport=transport, follow_redirects=True),
|
|
)
|
|
settings = Settings(
|
|
data_dir=tmp_path / "data",
|
|
remote_cache_gb=0.01,
|
|
resource_disk_reserve_gb=0.5,
|
|
)
|
|
settings.prepare()
|
|
cache = RemoteMediaCache(settings, _RemoteSource())
|
|
video = {
|
|
"id": "video-1",
|
|
"source_id": "source-1",
|
|
"source_key": "folder/video.mp4",
|
|
"fingerprint": "etag-1",
|
|
"size_bytes": len(payload),
|
|
}
|
|
|
|
first = cache.acquire(video)
|
|
assert first.path.read_bytes() == payload
|
|
first.close()
|
|
second = cache.acquire(video)
|
|
assert second.path == first.path
|
|
second.close()
|
|
assert len(requests) == 1
|
|
assert cache.status()["used_bytes"] == len(payload)
|
|
|
|
|
|
def test_remote_media_partial_download_resumes(tmp_path: Path, monkeypatch):
|
|
payload = b"abcdefghij"
|
|
observed_range = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
range_header = request.headers.get("range")
|
|
observed_range.append(range_header)
|
|
offset = int((range_header or "bytes=0-").removeprefix("bytes=").removesuffix("-"))
|
|
return httpx.Response(
|
|
206 if offset else 200,
|
|
headers={"Content-Range": f"bytes {offset}-{len(payload) - 1}/{len(payload)}"},
|
|
content=payload[offset:],
|
|
)
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
client_type = httpx.Client
|
|
monkeypatch.setattr(
|
|
"imagefind.remote_cache.httpx.Client",
|
|
lambda **kwargs: client_type(transport=transport, follow_redirects=True),
|
|
)
|
|
settings = Settings(data_dir=tmp_path / "data", resource_disk_reserve_gb=0.5)
|
|
settings.prepare()
|
|
cache = RemoteMediaCache(settings, _RemoteSource())
|
|
video = {
|
|
"id": "video-2",
|
|
"source_id": "source-1",
|
|
"source_key": "video.mp4",
|
|
"fingerprint": "etag-2",
|
|
"size_bytes": len(payload),
|
|
}
|
|
_key, _data, part, _metadata = cache._paths(video)
|
|
part.write_bytes(payload[:4])
|
|
|
|
lease = cache.acquire(video)
|
|
assert lease.path.read_bytes() == payload
|
|
lease.close()
|
|
assert observed_range == ["bytes=4-"]
|
|
|
|
|
|
def test_active_partial_is_pinned_against_manual_cleanup(tmp_path: Path, monkeypatch):
|
|
payload = b"active-download"
|
|
settings = Settings(data_dir=tmp_path / "data", resource_disk_reserve_gb=0.5)
|
|
settings.prepare()
|
|
cache = RemoteMediaCache(settings, _RemoteSource())
|
|
video = {
|
|
"id": "video-active",
|
|
"source_id": "source-1",
|
|
"source_key": "active.mp4",
|
|
"fingerprint": "etag-active",
|
|
"size_bytes": len(payload),
|
|
}
|
|
|
|
def download(_video, part, **_kwargs):
|
|
part.write_bytes(payload)
|
|
result = cache.cleanup(force=True)
|
|
assert result["skipped_active"] == 1
|
|
assert part.read_bytes() == payload
|
|
|
|
monkeypatch.setattr(cache, "_download", download)
|
|
lease = cache.acquire(video)
|
|
assert lease.path.read_bytes() == payload
|
|
lease.close()
|
|
|
|
|
|
def test_rclone_instances_are_isolated_and_removed_on_stop(tmp_path: Path, monkeypatch):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
|
|
manager = RcloneManager(settings)
|
|
commands = []
|
|
ports = iter((41001, 41002))
|
|
|
|
class Process:
|
|
next_pid = 9100
|
|
|
|
def __init__(self):
|
|
self.pid = Process.next_pid
|
|
Process.next_pid += 1
|
|
self.returncode = None
|
|
|
|
def poll(self):
|
|
return self.returncode
|
|
|
|
def terminate(self):
|
|
self.returncode = 0
|
|
|
|
def kill(self):
|
|
self.returncode = -9
|
|
|
|
def wait(self, timeout=None):
|
|
return self.returncode
|
|
|
|
class Connection:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def popen(command, **_kwargs):
|
|
commands.append(command)
|
|
return Process()
|
|
|
|
monkeypatch.setattr(manager, "_binary", lambda: "/usr/bin/rclone")
|
|
monkeypatch.setattr(manager, "_obscure", lambda value: f"obscured-{value}")
|
|
monkeypatch.setattr(manager, "_free_port", lambda: next(ports))
|
|
monkeypatch.setattr("imagefind.remote.subprocess.Popen", popen)
|
|
monkeypatch.setattr("imagefind.remote.socket.create_connection", lambda *_args, **_kwargs: Connection())
|
|
|
|
def source(source_id):
|
|
return {
|
|
"id": source_id,
|
|
"config": {
|
|
"driver": "alist",
|
|
"mode": "encrypted",
|
|
"base_url": "https://alist.test",
|
|
"root_path": "videos",
|
|
"username": "admin",
|
|
},
|
|
"secrets": {"password": "dav", "crypt_password": "crypt", "crypt_salt": "salt"},
|
|
}
|
|
|
|
first = manager.ensure(source("source-a"))
|
|
second = manager.ensure(source("source-b"))
|
|
first_cache = Path(commands[0][commands[0].index("--cache-dir") + 1])
|
|
second_cache = Path(commands[1][commands[1].index("--cache-dir") + 1])
|
|
assert first_cache != second_cache
|
|
assert first.instance_dir in first_cache.parents
|
|
assert second.instance_dir in second_cache.parents
|
|
assert "--read-only" in commands[0]
|
|
assert commands[0][commands[0].index("--vfs-cache-mode") + 1] == "off"
|
|
assert "--vfs-cache-max-size" not in commands[0]
|
|
assert commands[0][commands[0].index("--buffer-size") + 1] == "4M"
|
|
|
|
manager.stop_all()
|
|
assert not first.instance_dir.exists()
|
|
assert not second.instance_dir.exists()
|
|
|
|
|
|
def test_rclone_start_failure_does_not_leave_instance_files(tmp_path: Path, monkeypatch):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
|
|
manager = RcloneManager(settings)
|
|
monkeypatch.setattr(manager, "_binary", lambda: "/usr/bin/rclone")
|
|
monkeypatch.setattr(manager, "_obscure", lambda value: value)
|
|
monkeypatch.setattr(manager, "_free_port", lambda: 41003)
|
|
|
|
def fail_start(*_args, **_kwargs):
|
|
raise OSError("cannot start")
|
|
|
|
monkeypatch.setattr("imagefind.remote.subprocess.Popen", fail_start)
|
|
source = {
|
|
"id": "source-failed",
|
|
"config": {
|
|
"driver": "alist",
|
|
"mode": "encrypted",
|
|
"base_url": "https://alist.test",
|
|
"root_path": "",
|
|
"username": "admin",
|
|
},
|
|
"secrets": {"password": "dav", "crypt_password": "crypt", "crypt_salt": "salt"},
|
|
}
|
|
|
|
try:
|
|
manager.ensure(source)
|
|
except OSError as exc:
|
|
assert "cannot start" in str(exc)
|
|
else:
|
|
raise AssertionError("rclone startup failure must be surfaced")
|
|
assert not list((settings.rclone_dir / "instances" / "source-failed").glob("*"))
|
|
|
|
|
|
def test_rclone_orphan_is_reclaimed_after_three_idle_samples(tmp_path: Path, monkeypatch):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
reconcile = RcloneManager.reconcile_orphans
|
|
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
|
|
manager = RcloneManager(settings)
|
|
monkeypatch.setattr(RcloneManager, "reconcile_orphans", reconcile)
|
|
target_pid = os.getpid()
|
|
times = iter((1000.0, 1060.0, 1121.0, 1122.0))
|
|
reaped: list[int] = []
|
|
|
|
monkeypatch.setattr("imagefind.remote.time.monotonic", lambda: next(times))
|
|
monkeypatch.setattr(manager, "_process_identity", lambda pid: (1, 77) if pid == target_pid else None)
|
|
monkeypatch.setattr(manager, "_managed_command", lambda pid: pid == target_pid)
|
|
monkeypatch.setattr(manager, "_registry_matches", lambda _pid, _ticks: None)
|
|
monkeypatch.setattr(manager, "_io_bytes", lambda _pid: (10, 20))
|
|
monkeypatch.setattr(manager, "_has_established_connection", lambda _pid: False)
|
|
monkeypatch.setattr(manager, "_has_open_cache_file", lambda _pid: True)
|
|
monkeypatch.setattr(manager, "_terminate_pid", lambda pid: reaped.append(pid))
|
|
monkeypatch.setattr(manager, "_cleanup_registered_instance", lambda *_args: None)
|
|
|
|
assert manager.reconcile_orphans()["verifying"] == 1
|
|
assert manager.reconcile_orphans()["verifying"] == 1
|
|
result = manager.reconcile_orphans()
|
|
|
|
assert result == {"matched": 1, "reaped": 1, "verifying": 0, "unresolved": 0}
|
|
assert reaped == [target_pid]
|
|
assert manager.status()["reclaimed"] == 1
|
|
|
|
|
|
def test_rclone_cleanup_removes_only_stale_vfs_cache(tmp_path: Path, monkeypatch):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
|
|
manager = RcloneManager(settings)
|
|
stale = settings.rclone_dir / "instances" / "source-stale" / "instance-stale"
|
|
cache = stale / "cache" / "vfs" / "encrypted"
|
|
cache.mkdir(parents=True)
|
|
(cache / "movie.bin").write_bytes(b"cached" * 1024)
|
|
(stale / "instance.json").write_text('{"pid":999999,"start_ticks":1}', encoding="utf-8")
|
|
|
|
before = manager.status()
|
|
assert before["vfs_cache_bytes"] == 6 * 1024
|
|
assert before["reclaimable_cache_bytes"] == 6 * 1024
|
|
|
|
result = manager.cleanup_vfs_cache()
|
|
|
|
assert result == {"freed_bytes": 6 * 1024, "removed": 1, "skipped_active": 0}
|
|
assert not (stale / "cache").exists()
|
|
|
|
|
|
def test_rclone_accounts_for_and_cleans_legacy_shared_vfs_cache(tmp_path: Path, monkeypatch):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
monkeypatch.setattr(RcloneManager, "reconcile_orphans", lambda self: {})
|
|
manager = RcloneManager(settings)
|
|
legacy = settings.rclone_dir / "cache" / "vfs" / "crypt_source"
|
|
legacy.mkdir(parents=True)
|
|
(legacy / "old-video.bin").write_bytes(b"legacy" * 1024)
|
|
|
|
before = manager.status()
|
|
assert before["vfs_cache_bytes"] == 6 * 1024
|
|
assert before["reclaimable_cache_bytes"] == 6 * 1024
|
|
assert before["runtime_bytes"] == 0
|
|
assert before["instances"][0]["instance_id"] == "legacy-cache"
|
|
|
|
result = manager.cleanup_vfs_cache()
|
|
|
|
assert result == {"freed_bytes": 6 * 1024, "removed": 1, "skipped_active": 0}
|
|
assert not (settings.rclone_dir / "cache").exists()
|
|
|
|
|
|
def test_storage_usage_exposes_nested_cache_and_staging_details(tmp_path: Path):
|
|
settings = Settings(data_dir=tmp_path / "data")
|
|
settings.prepare()
|
|
(settings.models_dir / "model.bin").write_bytes(b"m" * 10)
|
|
(settings.remote_media_cache_dir / "video.media").write_bytes(b"r" * 20)
|
|
external = tmp_path / "openlist-stage"
|
|
external.mkdir()
|
|
(external / "encrypted.bin").write_bytes(b"e" * 30)
|
|
outside = tmp_path / "outside.bin"
|
|
outside.write_bytes(b"x" * 100)
|
|
(settings.remote_media_cache_dir / "ignored-link").symlink_to(outside)
|
|
|
|
class Connection:
|
|
def execute(self, sql, _params=()):
|
|
if "pg_database_size" in sql:
|
|
return type("Result", (), {"fetchone": lambda self: (1000,)})()
|
|
if "pg_total_relation_size" in sql:
|
|
return type("Result", (), {"fetchone": lambda self: (200,)})()
|
|
if "FROM sources" in sql:
|
|
rows = [
|
|
{
|
|
"id": "source-1",
|
|
"name": "OpenList",
|
|
"config_json": f'{{"openlist_local_staging_path":"{external}"}}',
|
|
}
|
|
]
|
|
return type("Result", (), {"fetchall": lambda self: rows})()
|
|
raise AssertionError(sql)
|
|
|
|
class Read:
|
|
def __enter__(self):
|
|
return Connection()
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
class Db:
|
|
def read(self):
|
|
return Read()
|
|
|
|
class Rclone:
|
|
def status(self):
|
|
return {
|
|
"vfs_cache_bytes": 40,
|
|
"active_cache_bytes": 5,
|
|
"reclaimable_cache_bytes": 35,
|
|
"runtime_bytes": 3,
|
|
"instances": [],
|
|
}
|
|
|
|
class Cache:
|
|
def status(self):
|
|
return {"active_bytes": 4, "evictable_bytes": 16, "entries": 1}
|
|
|
|
usage = StorageUsageService(Db(), settings, Rclone(), Cache()).usage(refresh=True)
|
|
|
|
assert usage["schema_version"] == 2
|
|
assert usage["categories"]["remote_cache"]["children"]["remote_media_cache"]["bytes"] == 20
|
|
assert usage["categories"]["remote_cache"]["children"]["rclone_vfs_cache"]["bytes"] == 40
|
|
openlist = usage["categories"]["upload_staging"]["children"]["openlist_local_staging"]
|
|
assert openlist["bytes"] == 30
|
|
assert openlist["available"] is True
|
|
assert usage["app_bytes"] == _path_bytes(settings.data_dir) + 1000 + 30
|
|
|
|
|
|
def _path_bytes(root: Path) -> int:
|
|
return sum(path.stat().st_size for path in root.rglob("*") if path.is_file() and not path.is_symlink())
|