From 3b6c46d2ee8c7136c49ce62c3c3a2f5020922c57 Mon Sep 17 00:00:00 2001 From: jcbl1 Date: Fri, 23 Oct 2026 09:28:57 +0000 Subject: [PATCH 1/2] perf(storage): lazy-load backends, stream chunk merges, offload sync IO - import storage backends on first use so local mode never loads aioboto3/botocore/aiohttp - reuse S3/WebDAV/OpenDAL instances while their config signature is unchanged - merge local chunks in 256KB blocks and hash incrementally - stream OpenDAL save/merge via open("wb") with a legacy fallback - move blocking local fs cleanup and share-dir walks into asyncio.to_thread - cover lazy imports, cache invalidation, streaming bounds and thread offload --- apps/base/tasks.py | 13 ++- core/storage/__init__.py | 93 +++++++++++++++--- core/storage/_base.py | 19 +++- core/storage/local.py | 148 ++++++++++++++++++---------- core/storage/opendal.py | 85 ++++++++++++++-- core/storage/s3.py | 20 ++++ core/storage/webdav.py | 14 +++ tests/test_cleanup_tasks.py | 35 +++++++ tests/test_local_storage_io.py | 126 ++++++++++++++++++++++++ tests/test_merge_chunks.py | 71 ++++++++++++++ tests/test_opendal_merge_stream.py | 152 +++++++++++++++++++++++++++++ tests/test_storage_lazy_cache.py | 129 ++++++++++++++++++++++++ 12 files changed, 830 insertions(+), 75 deletions(-) create mode 100644 tests/test_local_storage_io.py create mode 100644 tests/test_opendal_merge_stream.py create mode 100644 tests/test_storage_lazy_cache.py diff --git a/apps/base/tasks.py b/apps/base/tasks.py index c45505cd8..79168e985 100644 --- a/apps/base/tasks.py +++ b/apps/base/tasks.py @@ -23,16 +23,23 @@ from apps.base.local_share import is_local_ref from core.utils import get_now + +def _prune_empty_share_dirs(root: str) -> None: + """删除 share/data 下的空目录（阻塞操作，调用方需放入线程池）。""" + for current, dirs, files in os.walk(os.path.join(str(root), "share", "data")): + if not dirs and not files: + os.rmdir(current) + + async def delete_expire_files(): while True: try: await refresh_settings() file_storage: FileStorageInterface = storages[settings.file_storage]() # 遍历 share目录下的所有文件夹，删除空的文件夹，并判断父目录是否为空，如果为空也删除 + # 目录遍历是阻塞的文件系统调用，下放线程池避免卡住事件循环 if settings.file_storage == "local": - for root, dirs, files in os.walk(f"{data_root}/share/data"): - if not dirs and not files: - os.rmdir(root) + await asyncio.to_thread(_prune_empty_share_dirs, data_root) await ip_limit["error"].remove_expired_ip() await ip_limit["metadata"].remove_expired_ip() await ip_limit["upload"].remove_expired_ip() diff --git a/core/storage/__init__.py b/core/storage/__init__.py index 757e31532..c2494360b 100644 --- a/core/storage/__init__.py +++ b/core/storage/__init__.py @@ -2,7 +2,20 @@ Public surface is re-exported here — `from core.storage import X` keeps working for every historical name. + +Backends load lazily: importing this package only pulls in ``_base`` (the +framework-free contracts). Selecting ``local`` therefore never imports +aioboto3/botocore/aiohttp, which is the bulk of the resident memory on a +small VPS. ``storages[name]()`` keeps its historical shape; backends that +opt in via ``cache_signature()`` are also reused across requests while their +configuration is unchanged (previously every request rebuilt e.g. an +``aioboto3.Session``). """ +from __future__ import annotations + +from importlib import import_module +from typing import Callable + from core.storage._base import ( # noqa: F401 FileStorageInterface, S3_MIN_MULTIPART_PART_SIZE, @@ -10,16 +23,72 @@ StoredFile, build_attachment_headers, ) -from core.storage.local import SystemFileStorage # noqa: F401 -from core.storage.s3 import S3FileStorage # noqa: F401 -from core.storage.onedrive import OneDriveFileStorage # noqa: F401 -from core.storage.opendal import OpenDALFileStorage # noqa: F401 -from core.storage.webdav import WebDAVFileStorage # noqa: F401 - -storages = { - "local": SystemFileStorage, - "s3": S3FileStorage, - "onedrive": OneDriveFileStorage, - "opendal": OpenDALFileStorage, - "webdav": WebDAVFileStorage, + +# 历史导出名 -> (模块, 类名)。模块级 __getattr__ 与惰性工厂共用这张表。 +_LAZY_EXPORTS: dict[str, tuple[str, str]] = { + "SystemFileStorage": ("core.storage.local", "SystemFileStorage"), + "S3FileStorage": ("core.storage.s3", "S3FileStorage"), + "OneDriveFileStorage": ("core.storage.onedrive", "OneDriveFileStorage"), + "OpenDALFileStorage": ("core.storage.opendal", "OpenDALFileStorage"), + "WebDAVFileStorage": ("core.storage.webdav", "WebDAVFileStorage"), +} + + +def _load_class(module_name: str, attr: str): + return getattr(import_module(module_name), attr) + + +class _StorageFactory: + """Lazily import a backend class and reuse its instance while config holds. + + Call sites continue to use ``storages[name]()``. Backends exposing a + ``cache_signature()`` classmethod are constructed once per configuration + and reused; a changed signature rebuilds the instance. ``local`` and + ``onedrive`` deliberately omit it (see their module notes). + """ + + def __init__(self, module_name: str, attr: str): + self._module_name = module_name + self._attr = attr + self._cached: tuple[object, FileStorageInterface] | None = None + + @property + def storage_class(self): + return _load_class(self._module_name, self._attr) + + def __call__(self) -> FileStorageInterface: + cls = self.storage_class + signature_fn = getattr(cls, "cache_signature", None) + if signature_fn is None: + return cls() + signature = signature_fn() + if self._cached is not None and self._cached[0] == signature: + return self._cached[1] + instance = cls() + self._cached = (signature, instance) + return instance + + def __repr__(self) -> str: + return f"" + + +storages: dict[str, Callable[[], FileStorageInterface]] = { + "local": _StorageFactory("core.storage.local", "SystemFileStorage"), + "s3": _StorageFactory("core.storage.s3", "S3FileStorage"), + "onedrive": _StorageFactory("core.storage.onedrive", "OneDriveFileStorage"), + "opendal": _StorageFactory("core.storage.opendal", "OpenDALFileStorage"), + "webdav": _StorageFactory("core.storage.webdav", "WebDAVFileStorage"), } + + +def __getattr__(name: str): + target = _LAZY_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = _load_class(*target) + globals()[name] = value # 缓存解析结果，后续访问不再经过 __getattr__ + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(_LAZY_EXPORTS)) diff --git a/core/storage/_base.py b/core/storage/_base.py index 47ff99db8..ef2495524 100644 --- a/core/storage/_base.py +++ b/core/storage/_base.py @@ -86,6 +86,17 @@ def _get_chunk_record(chunk_records: dict, index: int): raise ValueError(f"分片{index}记录不存在") return chunk_record + def _verify_chunk_hash(self, index: int, chunk_record, current_hash: str) -> None: + """Compare a chunk digest with its record using the shared error wording. + + Streaming backends hash incrementally and call this once per chunk, so + the message stays byte-identical to the in-memory path. + """ + if current_hash != chunk_record.chunk_hash: + raise ValueError( + f"分片{index}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}" + ) + def _verify_and_hash_chunk( self, index: int, @@ -95,11 +106,9 @@ def _verify_and_hash_chunk( ) -> None: """Verify a chunk against its recorded hash, then stripe it into the whole-file digest. Raises the shared ValueError wording on mismatch.""" - current_hash = hashlib.sha256(chunk_data).hexdigest() - if current_hash != chunk_record.chunk_hash: - raise ValueError( - f"分片{index}哈希不匹配: 期望 {chunk_record.chunk_hash}, 实际 {current_hash}" - ) + self._verify_chunk_hash( + index, chunk_record, hashlib.sha256(chunk_data).hexdigest() + ) file_sha256.update(chunk_data) async def save_file( diff --git a/core/storage/local.py b/core/storage/local.py index 75942f5f3..df5a8b23c 100644 --- a/core/storage/local.py +++ b/core/storage/local.py @@ -43,6 +43,49 @@ def _save(self, file, save_path): f.write(chunk) chunk = file.read(self.chunk_size) + def _prepare_save_path(self, directory: str, filename: str) -> Path: + """Resolve + create the parent directory (blocking; run in a thread).""" + safe_save_path = self._resolve_safe_path( + f"{directory}/{filename}" if directory not in {"", "."} else filename + ) + safe_save_path.parent.mkdir(parents=True, exist_ok=True) + return safe_save_path + + def _prepare_chunk_path( + self, save_path: str, upload_id: str, chunk_index: int + ) -> Path: + """Validate the target path and materialise the chunk directory.""" + self._resolve_safe_path(save_path) + chunk_path = self._resolve_safe_path( + str(Path(save_path).parent / "chunks" / upload_id / f"{chunk_index}.part") + ) + chunk_path.parent.mkdir(parents=True, exist_ok=True) + return chunk_path + + @staticmethod + def _unlink_if_exists(path: Path) -> None: + try: + path.unlink() + except FileNotFoundError: + pass + + def _delete_file_sync(self, relative_path: str) -> None: + self._unlink_if_exists(self._resolve_safe_path(relative_path)) + + def _stat_file_sync(self, relative_path: str) -> tuple[Path, int | None] | None: + """Return (path, size); None when the file is gone. + + A stat that fails for another reason still yields the path with an + unknown size, matching the old "omit Content-Length" behaviour. + """ + file_path = self._resolve_safe_path(relative_path) + try: + return file_path, file_path.stat().st_size + except FileNotFoundError: + return None + except OSError: + return file_path, None + async def save_file( self, stream: BinaryIO, save_path: str, content_type: Optional[str] = None ): @@ -50,32 +93,28 @@ async def save_file( directory = str(path_obj.parent).replace("\\", "/").lstrip("/") # 提取原始文件名并进行清理 filename = await sanitize_filename(path_obj.name) - # 构建安全的完整保存路径 - safe_save_path = self._resolve_safe_path(f"{directory}/{filename}" if directory not in {"", "."} else filename) - # 确保目录存在 - if not safe_save_path.parent.exists(): - safe_save_path.parent.mkdir(parents=True) + # 构建安全的完整保存路径（解析路径与建目录都是阻塞系统调用） + safe_save_path = await asyncio.to_thread( + self._prepare_save_path, directory, filename + ) await asyncio.to_thread(self._save, stream, safe_save_path) async def delete_file(self, file_code: StoredFile): - save_path = self._resolve_safe_path(file_code.get_file_path()) - if save_path.exists(): - save_path.unlink() + await asyncio.to_thread(self._delete_file_sync, file_code.get_file_path()) async def get_file_url(self, file_code: StoredFile): return await get_file_url(file_code.code) async def get_file_response(self, file_code: StoredFile): - file_path = self._resolve_safe_path(file_code.get_file_path()) - if not file_path.exists(): + stat_result = await asyncio.to_thread( + self._stat_file_sync, file_code.get_file_path() + ) + if stat_result is None: raise StorageError(status_code=404, detail="文件已过期删除") + file_path, size = stat_result filename = f"{file_code.prefix}{file_code.suffix}" - try: - headers = build_attachment_headers(filename, file_path.stat().st_size) - except OSError: - # 文件大小不可得时省略 Content-Length - headers = build_attachment_headers(filename) - + headers = build_attachment_headers(filename, size) + return StoredDownload( filename=filename, headers=headers, @@ -91,40 +130,39 @@ async def save_chunk(self, upload_id: str, chunk_index: int, chunk_data: bytes, :param chunk_hash: 分片哈希值 :param save_path: 文件保存路径 """ - # 先校验目标文件路径合法，再将分片落到同级 chunks 目录。 - self._resolve_safe_path(save_path) - chunk_path = self._resolve_safe_path( - str(Path(save_path).parent / "chunks" / upload_id / f"{chunk_index}.part") + chunk_path = await asyncio.to_thread( + self._prepare_chunk_path, save_path, upload_id, chunk_index ) - if not chunk_path.parent.exists(): - chunk_path.parent.mkdir(parents=True, exist_ok=True) # 使用临时文件写入，确保原子性 temp_path = chunk_path.with_suffix('.tmp') try: async with aiofiles.open(temp_path, "wb") as f: await f.write(chunk_data) # 原子重命名 - temp_path.rename(chunk_path) + await asyncio.to_thread(temp_path.rename, chunk_path) except Exception as e: - if temp_path.exists(): - temp_path.unlink() + await asyncio.to_thread(self._unlink_if_exists, temp_path) raise e async def merge_chunks(self, upload_id: str, total_chunks: int, chunk_size: int, save_path: str, chunk_records: dict) -> tuple[str, str]: """ 合并本地文件系统的分片文件并返回文件路径和完整哈希值 + + 按固定块（self.chunk_size）流式读写：内存峰值与分片大小无关， + 只与块大小有关（旧实现会把整片 read() 进内存）。 :param upload_id: 上传会话ID :param chunk_info: 分片信息 :param save_path: 文件保存路径 :return: (文件路径, 文件哈希值) """ - output_path = self._resolve_safe_path(save_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - chunk_base_dir = self._resolve_safe_path( - str(Path(save_path).parent / "chunks" / upload_id) + output_path = await asyncio.to_thread(self._resolve_safe_path, save_path) + await asyncio.to_thread(output_path.parent.mkdir, parents=True, exist_ok=True) + chunk_base_dir = await asyncio.to_thread( + self._resolve_safe_path, + str(Path(save_path).parent / "chunks" / upload_id), ) file_sha256 = hashlib.sha256() - + # 使用临时文件写入，确保原子性 temp_output = output_path.with_suffix('.merging') try: @@ -133,26 +171,27 @@ async def merge_chunks(self, upload_id: str, total_chunks: int, chunk_size: int, # 获取分片记录 chunk_record = self._get_chunk_record(chunk_records, i) chunk_path = chunk_base_dir / f"{i}.part" - if not chunk_path.exists(): - raise ValueError(f"分片{i}文件不存在") - async with aiofiles.open(chunk_path, "rb") as in_file: - chunk_data = await in_file.read() - self._verify_and_hash_chunk(i, chunk_record, chunk_data, file_sha256) - await out_file.write(chunk_data) + chunk_sha256 = hashlib.sha256() + try: + async with aiofiles.open(chunk_path, "rb") as in_file: + while True: + block = await in_file.read(self.chunk_size) + if not block: + break + chunk_sha256.update(block) + file_sha256.update(block) + await out_file.write(block) + except FileNotFoundError as exc: + raise ValueError(f"分片{i}文件不存在") from exc + self._verify_chunk_hash(i, chunk_record, chunk_sha256.hexdigest()) # 原子重命名 - temp_output.rename(output_path) + await asyncio.to_thread(temp_output.rename, output_path) except Exception as e: - if temp_output.exists(): - temp_output.unlink() + await asyncio.to_thread(self._unlink_if_exists, temp_output) raise e return str(output_path), file_sha256.hexdigest() - async def clean_chunks(self, upload_id: str, save_path: str): - """ - 清理本地文件系统的临时分片文件 - :param upload_id: 上传会话ID - :param save_path: 文件保存路径 - """ + def _clean_chunks_sync(self, upload_id: str, save_path: str) -> None: chunk_dir = self._resolve_safe_path( str(Path(save_path).parent / "chunks" / upload_id) ) @@ -169,16 +208,27 @@ async def clean_chunks(self, upload_id: str, save_path: str): except Exception as e: logger.warning(f"清理 chunks 父目录失败: {e}") - async def file_exists(self, save_path: str) -> bool: + async def clean_chunks(self, upload_id: str, save_path: str): """ - 检查文件是否存在于本地文件系统 - :param save_path: 文件路径 - :return: 文件是否存在 + 清理本地文件系统的临时分片文件 + :param upload_id: 上传会话ID + :param save_path: 文件保存路径 """ + await asyncio.to_thread(self._clean_chunks_sync, upload_id, save_path) + + def _file_exists_sync(self, save_path: str) -> bool: try: file_path = self._resolve_safe_path(save_path) except ValueError: return False return file_path.exists() + async def file_exists(self, save_path: str) -> bool: + """ + 检查文件是否存在于本地文件系统 + :param save_path: 文件路径 + :return: 文件是否存在 + """ + return await asyncio.to_thread(self._file_exists_sync, save_path) + diff --git a/core/storage/opendal.py b/core/storage/opendal.py index bdcb34d46..81946ec98 100644 --- a/core/storage/opendal.py +++ b/core/storage/opendal.py @@ -20,6 +20,19 @@ class OpenDALFileStorage(FileStorageInterface): + # 供 core.storage._StorageFactory 复用实例。AsyncOperator 设计上可共享， + # 签名跟随 opendal_scheme 及当前 scheme 的全部 opendal__* 配置。 + @classmethod + def cache_signature(cls) -> tuple: + scheme = getattr(settings, "opendal_scheme", "") + prefix = f"opendal_{scheme}_" + items = { + key: value + for key, value in settings.items() + if key == "opendal_scheme" or key.startswith(prefix) + } + return (repr(sorted(items.items())),) + def __init__(self): try: import opendal @@ -35,12 +48,45 @@ def __init__(self): settings.opendal_scheme, **service_settings ) + async def _open_writer(self, save_path: str): + """Open a streaming writer, or None when the installed SDK is too old.""" + try: + return await self.operator.open(save_path, "wb") + except (AttributeError, NotImplementedError, TypeError) as exc: + logger.warning("OpenDAL 不支持流式写入，回退到整文件写入: %s", exc) + return None + + async def _discard_partial_write(self, writer, save_path: str) -> None: + """关闭流式 writer 并尽力删除已部分写入的对象。""" + try: + await writer.close() + finally: + try: + await self.operator.delete(save_path) + except Exception as cleanup_error: + logger.warning( + "清理 OpenDAL 写入失败残留失败 %s: %s", save_path, cleanup_error + ) + async def save_file( self, stream: BinaryIO, save_path: str, content_type: Optional[str] = None ): - # 使用 asyncio.to_thread 避免阻塞事件循环 - content = await asyncio.to_thread(stream.read) - await self.operator.write(save_path, content) + writer = await self._open_writer(save_path) + if writer is None: + # 旧版 SDK 回退：整文件读入后一次写入（内存 = 文件大小） + content = await asyncio.to_thread(stream.read) + await self.operator.write(save_path, content) + return + try: + while True: + block = await asyncio.to_thread(stream.read, 256 * 1024) + if not block: + break + await writer.write(block) + except Exception: + await self._discard_partial_write(writer, save_path) + raise + await writer.close() async def delete_file(self, file_code: StoredFile): await self.operator.delete(file_code.get_file_path()) @@ -102,11 +148,38 @@ async def save_chunk(self, upload_id: str, chunk_index: int, chunk_data: bytes, await self.operator.write(chunk_path, chunk_data) async def merge_chunks(self, upload_id: str, total_chunks: int, chunk_size: int, save_path: str, chunk_records: dict) -> tuple[str, str]: - """合并 OpenDAL 存储上的分片文件，使用临时文件避免内存问题""" + """合并 OpenDAL 存储上的分片文件。 + + 首选流式 writer：内存上界 = 单个分片（operator.read 无法按块读）， + 不再把整个合并结果 read() 进内存。旧版 SDK 不支持 open("wb") 时 + 回退到临时文件 + 整文件写入，保留原有行为。 + """ file_sha256 = hashlib.sha256() chunk_dir = str(Path(save_path).parent / "chunks" / upload_id) - # 使用临时文件存储合并数据 + writer = await self._open_writer(save_path) + if writer is not None: + try: + for i in range(total_chunks): + chunk_path = f"{chunk_dir}/{i}.part" + chunk_record = self._get_chunk_record(chunk_records, i) + + try: + chunk_data = await self.operator.read(chunk_path) + except Exception as e: + raise ValueError(f"分片{i}文件不存在: {e}") + + self._verify_and_hash_chunk(i, chunk_record, chunk_data, file_sha256) + await writer.write(chunk_data) + del chunk_data # 释放分片内存 + except Exception: + # 失败时目标对象可能已被部分写入，尽力清理后再抛出。 + await self._discard_partial_write(writer, save_path) + raise + await writer.close() + return save_path, file_sha256.hexdigest() + + # 旧版 SDK 回退路径：临时文件避免重复网络读，最后整文件写入。 with tempfile.NamedTemporaryFile(delete=False) as temp_file: temp_path = temp_file.name @@ -133,7 +206,7 @@ async def merge_chunks(self, upload_id: str, total_chunks: int, chunk_size: int, # 清理临时文件 if os.path.exists(temp_path): os.unlink(temp_path) - + return save_path, file_sha256.hexdigest() async def clean_chunks(self, upload_id: str, save_path: str): diff --git a/core/storage/s3.py b/core/storage/s3.py index 2dd25eacc..8c96c2f5e 100644 --- a/core/storage/s3.py +++ b/core/storage/s3.py @@ -23,6 +23,26 @@ class S3FileStorage(FileStorageInterface): + # 供 core.storage._StorageFactory 复用实例；只列影响构造/签名的配置项。 + # 配置变更时签名变化，工厂会重建实例（见到新 endpoint/bucket 等）。 + @classmethod + def cache_signature(cls) -> tuple[str, ...]: + return tuple( + repr(getattr(settings, key, None)) + for key in ( + "s3_access_key_id", + "s3_secret_access_key", + "s3_bucket_name", + "s3_hostname", + "s3_region_name", + "s3_signature_version", + "s3_endpoint_url", + "aws_session_token", + "s3_addressing_style", + "s3_proxy", + ) + ) + def __init__(self): self.access_key_id = settings.s3_access_key_id self.secret_access_key = settings.s3_secret_access_key diff --git a/core/storage/webdav.py b/core/storage/webdav.py index 710f23ca0..4f32a0ba2 100644 --- a/core/storage/webdav.py +++ b/core/storage/webdav.py @@ -25,6 +25,20 @@ class WebDAVFileStorage(FileStorageInterface): + # 供 core.storage._StorageFactory 复用实例；只列影响构造结果的配置项。 + @classmethod + def cache_signature(cls) -> tuple[str, ...]: + return tuple( + repr(getattr(settings, key, None)) + for key in ( + "webdav_url", + "webdav_username", + "webdav_password", + "webdav_root_path", + "webdav_proxy", + ) + ) + def __init__(self): if not hasattr(self, "_initialized"): self.base_url = settings.webdav_url.rstrip("/") + "/" diff --git a/tests/test_cleanup_tasks.py b/tests/test_cleanup_tasks.py index 745aa5cf5..04c9d18eb 100644 --- a/tests/test_cleanup_tasks.py +++ b/tests/test_cleanup_tasks.py @@ -6,6 +6,8 @@ is isolated into a temp dir by patching data_root in both consumers. """ import asyncio +import os +import threading import unittest from pathlib import Path from tempfile import TemporaryDirectory @@ -140,6 +142,39 @@ async def _scenario(self): await close_db() +class PruneEmptyShareDirsTests(_OneRoundMixin, unittest.TestCase): + def test_empty_dirs_pruned_in_worker_thread(self): + asyncio.run(self._scenario()) + + async def _scenario(self): + with TemporaryDirectory() as tmpdir: + await init_memory_db() + try: + empty_dir = Path(tmpdir) / "share/data/2026/01/01/empty" + empty_dir.mkdir(parents=True) + kept_dir = Path(tmpdir) / "share/data/2026/01/01/kept" + kept_dir.mkdir(parents=True) + (kept_dir / "f.bin").write_bytes(b"x") + + recorded = {} + main_thread = threading.current_thread() + real_walk = os.walk + + def spy_walk(path, *args, **kwargs): + recorded["thread"] = threading.current_thread() + yield from real_walk(path, *args, **kwargs) + + with patch("apps.base.tasks.os.walk", spy_walk): + await self.run_one_round(delete_expire_files, tmpdir) + + self.assertTrue(recorded) + self.assertIsNot(recorded["thread"], main_thread, "目录遍历必须下放线程") + self.assertFalse(empty_dir.exists(), "空目录应被清理") + self.assertTrue(kept_dir.exists(), "非空目录必须保留") + finally: + await close_db() + + class CleanExpiredPresignSessionsTests(_OneRoundMixin, unittest.TestCase): def test_expired_direct_session_file_and_rows_cleaned(self): asyncio.run(self._scenario()) diff --git a/tests/test_local_storage_io.py b/tests/test_local_storage_io.py new file mode 100644 index 000000000..ec3f7aea8 --- /dev/null +++ b/tests/test_local_storage_io.py @@ -0,0 +1,126 @@ +"""本地存储的阻塞文件操作必须下放线程池。 + +单核 VPS 上，事件循环里直接 unlink/rmtree/stat 会卡住所有并发请求； +这里用线程记录器钉住这些边界（测试进程里 asyncio.run 的主线程名为 +MainThread，to_thread 运行在 asyncio_* 工作线程）。 +""" +import shutil +import threading +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +import pytest + +from core.errors import StorageError +from core.storage import StoredFile, SystemFileStorage + + +def _file_code(save_path: str) -> StoredFile: + return StoredFile( + file_path=save_path.rsplit("/", 1)[0], + uuid_file_name=save_path.rsplit("/", 1)[1], + ) + + +async def test_delete_file_runs_in_worker_thread(): + with TemporaryDirectory() as tmpdir: + with patch("core.storage.local.data_root", Path(tmpdir)): + storage = SystemFileStorage() + target = Path(tmpdir) / "share/data/2026/01/01/uuid/f.bin" + target.parent.mkdir(parents=True) + target.write_bytes(b"payload") + + main_thread = threading.current_thread() + recorded = {} + real_unlink = Path.unlink + + def spy_unlink(self, *args, **kwargs): + recorded["thread"] = threading.current_thread() + return real_unlink(self, *args, **kwargs) + + with patch.object(Path, "unlink", spy_unlink): + await storage.delete_file( + _file_code("share/data/2026/01/01/uuid/f.bin") + ) + + assert recorded["thread"] is not main_thread + assert not target.exists() + + +async def test_clean_chunks_runs_in_worker_thread(): + with TemporaryDirectory() as tmpdir: + with patch("core.storage.local.data_root", Path(tmpdir)): + storage = SystemFileStorage() + chunk_dir = Path(tmpdir) / "share/data/2026/01/01/uuid/chunks/upload-1" + chunk_dir.mkdir(parents=True) + (chunk_dir / "0.part").write_bytes(b"x") + + main_thread = threading.current_thread() + recorded = {} + real_rmtree = shutil.rmtree + + def spy_rmtree(path, *args, **kwargs): + recorded["thread"] = threading.current_thread() + return real_rmtree(path, *args, **kwargs) + + with patch("core.storage.local.shutil.rmtree", spy_rmtree): + await storage.clean_chunks( + "upload-1", "share/data/2026/01/01/uuid/f.bin" + ) + + assert recorded["thread"] is not main_thread + assert not chunk_dir.exists() + + +async def test_missing_download_stats_in_worker_thread(): + with TemporaryDirectory() as tmpdir: + with patch("core.storage.local.data_root", Path(tmpdir)): + storage = SystemFileStorage() + (Path(tmpdir) / "share/data/2026/01/01/uuid").mkdir(parents=True) + + main_thread = threading.current_thread() + recorded = {} + real_stat = Path.stat + + def spy_stat(self, *args, **kwargs): + recorded["thread"] = threading.current_thread() + return real_stat(self, *args, **kwargs) + + with patch.object(Path, "stat", spy_stat): + with pytest.raises(StorageError) as exc_info: + await storage.get_file_response( + _file_code("share/data/2026/01/01/uuid/ghost.bin") + ) + + assert exc_info.value.status_code == 404 + assert recorded["thread"] is not main_thread + + +async def test_save_chunk_prepares_directory_in_worker_thread(): + with TemporaryDirectory() as tmpdir: + with patch("core.storage.local.data_root", Path(tmpdir)): + storage = SystemFileStorage() + main_thread = threading.current_thread() + recorded = {} + real_prepare = SystemFileStorage._prepare_chunk_path + + def spy_prepare(self, save_path, upload_id, chunk_index): + recorded["thread"] = threading.current_thread() + return real_prepare(self, save_path, upload_id, chunk_index) + + with patch.object(SystemFileStorage, "_prepare_chunk_path", spy_prepare): + await storage.save_chunk( + "upload-1", + 0, + b"payload", + "unused", + "share/data/2026/01/01/uuid/f.bin", + ) + + assert recorded["thread"] is not main_thread + chunk_path = ( + Path(tmpdir) + / "share/data/2026/01/01/uuid/chunks/upload-1/0.part" + ) + assert chunk_path.read_bytes() == b"payload" diff --git a/tests/test_merge_chunks.py b/tests/test_merge_chunks.py index 03a7e2256..d2c507082 100644 --- a/tests/test_merge_chunks.py +++ b/tests/test_merge_chunks.py @@ -12,6 +12,7 @@ import unittest from pathlib import Path from tempfile import TemporaryDirectory +from types import SimpleNamespace from unittest.mock import patch from tests.helpers import SettingsOverrideMixin, close_db, init_memory_db @@ -114,5 +115,75 @@ async def records(): await close_db() +class MergeChunksStreamingTests(unittest.TestCase): + """合并必须按块读取：峰值内存与分片大小无关，且哈希仍逐片校验。""" + + def test_reads_chunks_in_bounded_blocks(self): + asyncio.run(self._scenario()) + + async def _scenario(self): + import core.storage.local as local_module + + with TemporaryDirectory() as tmpdir: + with patch("core.storage.local.data_root", Path(tmpdir)): + storage = SystemFileStorage() + storage.chunk_size = 3 + save_path = "share/data/2026/01/01/uuid-stream/merged.bin" + chunk_dir = ( + Path(tmpdir) / "share/data/2026/01/01/uuid-stream/chunks/uid-stream" + ) + chunk_dir.mkdir(parents=True) + payloads = [b"abcdefghij", b"klm"] + for index, data in enumerate(payloads): + (chunk_dir / f"{index}.part").write_bytes(data) + + records = { + index: SimpleNamespace(chunk_hash=_sha(data)) + for index, data in enumerate(payloads) + } + + read_sizes = [] + real_open = local_module.aiofiles.open + + class RecordingFile: + def __init__(self, inner): + self._inner = inner + + async def read(self, *args): + read_sizes.append(args) + return await self._inner.read(*args) + + class RecordingContext: + def __init__(self, ctx): + self._ctx = ctx + + async def __aenter__(self): + return RecordingFile(await self._ctx.__aenter__()) + + async def __aexit__(self, *exc_info): + return await self._ctx.__aexit__(*exc_info) + + def fake_open(path, mode="r", *args, **kwargs): + ctx = real_open(path, mode, *args, **kwargs) + if "b" in mode and "r" in mode: + return RecordingContext(ctx) + return ctx + + with patch.object(local_module.aiofiles, "open", fake_open): + out_path, digest = await storage.merge_chunks( + "uid-stream", + len(payloads), + storage.chunk_size, + save_path, + records, + ) + + merged = b"".join(payloads) + self.assertEqual(Path(out_path).read_bytes(), merged) + self.assertEqual(digest, _sha(merged)) + # 每 3 字节一次读取：10 字节 -> 5 次（含 EOF），3 字节 -> 2 次 + self.assertEqual(read_sizes, [(3,)] * 7) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_opendal_merge_stream.py b/tests/test_opendal_merge_stream.py new file mode 100644 index 000000000..0a4da8e4c --- /dev/null +++ b/tests/test_opendal_merge_stream.py @@ -0,0 +1,152 @@ +"""OpenDAL 合并/上传的流式化契约（不依赖 opendal SDK）。 + +`core.storage.opendal` 只在构造函数里 import SDK，因此可以用 fake operator +钉住：优先走 open("wb") 流式写入；版本过旧时才回退整文件写入。 +""" +import hashlib +import io +from types import SimpleNamespace + +import pytest + +from core.storage import OpenDALFileStorage + + +def _sha(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +class FakeWriter: + def __init__(self, path, store, events): + self.path = path + self.store = store + self.events = events + + async def write(self, data): + self.events.append(("write", len(data))) + self.store[self.path] = self.store.get(self.path, b"") + data + return len(data) + + async def close(self): + self.events.append(("close",)) + + +class FakeOperator: + def __init__(self, files=None, support_open=True): + self.store = dict(files or {}) + self.support_open = support_open + self.events = [] + + async def open(self, path, mode): + if not self.support_open: + raise AttributeError("open") + assert mode == "wb" + return FakeWriter(path, self.store, self.events) + + async def read(self, path): + try: + return self.store[path] + except KeyError as exc: + raise FileNotFoundError(path) from exc + + async def write(self, path, data): + self.events.append(("fallback_write", len(data))) + self.store[path] = data + + async def delete(self, path): + self.events.append(("delete", path)) + self.store.pop(path, None) + + +def _storage(operator) -> OpenDALFileStorage: + storage = OpenDALFileStorage.__new__(OpenDALFileStorage) + storage.operator = operator + return storage + + +def _chunk_files(chunks): + base = "share/data/2026/01/01/uuid/chunks/upload-1" + return {f"{base}/{index}.part": data for index, data in chunks.items()} + + +def _records(chunks): + return { + index: SimpleNamespace(chunk_hash=_sha(data)) + for index, data in chunks.items() + } + + +async def test_merge_uses_streaming_writer(): + chunks = {0: b"aaaa", 1: b"bbbbbb"} + operator = FakeOperator(_chunk_files(chunks)) + storage = _storage(operator) + + save_path = "share/data/2026/01/01/uuid/merged.bin" + out, digest = await storage.merge_chunks( + "upload-1", 2, 4, save_path, _records(chunks) + ) + + merged = b"aaaa" + b"bbbbbb" + assert out == save_path + assert digest == _sha(merged) + assert operator.store[save_path] == merged + assert ("fallback_write", len(merged)) not in operator.events + assert [event for event in operator.events if event[0] == "write"] == [ + ("write", 4), + ("write", 6), + ] + assert ("close",) in operator.events + + +async def test_merge_hash_mismatch_removes_partial_object(): + chunks = {0: b"aaaa"} + operator = FakeOperator(_chunk_files(chunks)) + storage = _storage(operator) + records = {0: SimpleNamespace(chunk_hash=_sha(b"wrong"))} + + save_path = "share/data/2026/01/01/uuid/merged.bin" + with pytest.raises(ValueError, match="哈希不匹配"): + await storage.merge_chunks("upload-1", 1, 4, save_path, records) + + assert save_path not in operator.store + assert ("delete", save_path) in operator.events + + +async def test_merge_falls_back_when_open_unsupported(): + chunks = {0: b"aaaa", 1: b"bbbb"} + operator = FakeOperator(_chunk_files(chunks), support_open=False) + storage = _storage(operator) + + save_path = "share/data/2026/01/01/uuid/merged.bin" + _, digest = await storage.merge_chunks( + "upload-1", 2, 4, save_path, _records(chunks) + ) + + merged = b"aaaabbbb" + assert digest == _sha(merged) + assert operator.store[save_path] == merged + assert ("fallback_write", len(merged)) in operator.events + + +async def test_save_file_streams_in_bounded_blocks(): + payload = b"x" * (600 * 1024) + operator = FakeOperator() + storage = _storage(operator) + + await storage.save_file(io.BytesIO(payload), "share/data/big.bin") + + assert operator.store["share/data/big.bin"] == payload + writes = [event[1] for event in operator.events if event[0] == "write"] + assert len(writes) >= 2 + assert all(size <= 256 * 1024 for size in writes) + + +async def test_save_file_falls_back_when_open_unsupported(): + payload = b"y" * 1024 + operator = FakeOperator(support_open=False) + storage = _storage(operator) + + await storage.save_file(io.BytesIO(payload), "share/data/fallback.bin") + + assert operator.store["share/data/fallback.bin"] == payload + assert ("fallback_write", len(payload)) in operator.events diff --git a/tests/test_storage_lazy_cache.py b/tests/test_storage_lazy_cache.py new file mode 100644 index 000000000..227026be1 --- /dev/null +++ b/tests/test_storage_lazy_cache.py @@ -0,0 +1,129 @@ +"""Storage package lazy-loading and instance-reuse contracts. + +Importing the app must not drag in the heavy aioboto3/botocore/aiohttp +stack in local-storage deployments, and backends that opt into +``cache_signature()`` must be reused across requests until their +configuration changes. + +These run in a subprocess where needed because pytest itself (and other +tests) may already have imported the heavy modules into ``sys.modules``. +""" +import os +import subprocess +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +import core.storage as storage_pkg + +ROOT = Path(__file__).resolve().parents[1] +HEAVY_MODULES = ("aioboto3", "botocore", "aiohttp", "office365", "msal", "opendal") + + +def _run_python(code: str) -> subprocess.CompletedProcess: + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + dict.fromkeys([str(ROOT)] + [p for p in sys.path if p]) + ) + return subprocess.run( + [sys.executable, "-c", code], + cwd=ROOT, + env=env, + capture_output=True, + text=True, + ) + + +def test_importing_app_does_not_load_heavy_backends(): + code = ( + "import sys\n" + "import main\n" + f"heavy = [m for m in {HEAVY_MODULES!r} if m in sys.modules]\n" + "assert not heavy, heavy\n" + "print('ok')\n" + ) + result = _run_python(code) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "ok" + + +def test_importing_package_defers_backend_modules(): + code = ( + "import sys\n" + "import core.storage\n" + "assert 'core.storage.s3' not in sys.modules\n" + "assert 'core.storage.local' not in sys.modules\n" + f"heavy = [m for m in {HEAVY_MODULES!r} if m in sys.modules]\n" + "assert not heavy, heavy\n" + "from core.storage import S3FileStorage\n" + "assert 'core.storage.s3' in sys.modules\n" + "assert 'aioboto3' in sys.modules\n" + "print('ok')\n" + ) + result = _run_python(code) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "ok" + + +def test_registry_has_all_historical_backends(): + assert set(storage_pkg.storages) == { + "local", + "s3", + "onedrive", + "opendal", + "webdav", + } + + +def test_unknown_lazy_export_raises_attribute_error(): + with pytest.raises(AttributeError): + getattr(storage_pkg, "NoSuchStorage") + + +def test_local_backend_is_not_reused(): + """local 构造极轻，且测试/运行时会 patch data_root——必须每次新实例。""" + factory = storage_pkg.storages["local"] + assert factory() is not factory() + + +def test_factory_reuses_instance_until_config_changes(monkeypatch): + import core.storage.s3 as s3_module + from core.settings import settings + + created = [] + + class FakeS3: + def __init__(self): + created.append(self) + + @classmethod + def cache_signature(cls): + return (repr(getattr(settings, "s3_endpoint_url", None)),) + + monkeypatch.setattr(s3_module, "S3FileStorage", FakeS3) + factory = storage_pkg.storages["s3"] + monkeypatch.setattr(factory, "_cached", None) + + first = factory() + second = factory() + assert first is second + assert len(created) == 1 + + monkeypatch.setattr(settings, "s3_endpoint_url", "https://changed.example") + third = factory() + assert third is not first + assert len(created) == 2 + + +def test_patch_dict_still_intercepts_registry(): + """历史测试通过 patch.dict(views.storages, {...}) 注入 FakeStorage。""" + + class FakeStorage: + pass + + with patch.dict(storage_pkg.storages, {"local": FakeStorage}): + assert isinstance(storage_pkg.storages["local"](), FakeStorage) + + assert not isinstance(storage_pkg.storages["local"](), FakeStorage) From 42e9d1611261f9c3ad25f2cc62f67cc0c86c48ba Mon Sep 17 00:00:00 2001 From: jcbl1 Date: Fri, 23 Oct 2026 09:52:40 +0000 Subject: [PATCH 2/2] perf(tasks): batch cleanup scans and skip text payloads delete_expire_files, clean_incomplete_uploads and clean_expired_presign_sessions now walk rows with an id cursor in batches of 200, loading only the columns they need (text payloads are never selected). The cursor also keeps a failing delete from retrying the same batch forever. Covered by tests/test_cleanup_batching.py: >1 batch is fully drained and a failing row does not block progress. --- apps/base/tasks.py | 154 +++++++++++++++++++++------------ tests/test_cleanup_batching.py | 122 ++++++++++++++++++++++++++ 2 files changed, 220 insertions(+), 56 deletions(-) create mode 100644 tests/test_cleanup_batching.py diff --git a/apps/base/tasks.py b/apps/base/tasks.py index 79168e985..deb9e50cb 100644 --- a/apps/base/tasks.py +++ b/apps/base/tasks.py @@ -19,10 +19,12 @@ from core.logger import logger from core.settings import settings, data_root from core.storage import FileStorageInterface, StoredFile, storages -from apps.base.services import stored_file_of from apps.base.local_share import is_local_ref from core.utils import get_now +# 清理任务单批处理量：限制单次 ORM 实例数量与 text 大字段峰值。 +CLEANUP_BATCH_SIZE = 200 + def _prune_empty_share_dirs(root: str) -> None: """删除 share/data 下的空目录（阻塞操作，调用方需放入线程池）。""" @@ -43,20 +45,39 @@ async def delete_expire_files(): await ip_limit["error"].remove_expired_ip() await ip_limit["metadata"].remove_expired_ip() await ip_limit["upload"].remove_expired_ip() - await StorageReservation.filter(expires_at__lte=await get_now()).delete() - expire_data = await FileCodes.filter( - Q(expired_at__lt=await get_now()) | Q(expired_count=0) - ).all() - for exp in expire_data: - try: - if not is_local_ref(exp): - await file_storage.delete_file(stored_file_of(exp)) - except Exception as e: - logger.error(f"删除过期文件失败 code={exp.code}: {e}") - try: - await exp.delete() - except Exception as e: - logger.error(f"删除记录失败 code={exp.code}: {e}") + now = await get_now() + await StorageReservation.filter(expires_at__lte=now).delete() + # 游标分批：只取必要字段（不含 text），避免一次性载入全部过期记录； + # 用 id 游标而非纯 limit，单条删除失败时也不会在同一批里反复重试卡死。 + last_id = 0 + while True: + expire_data = await FileCodes.filter( + (Q(expired_at__lt=now) | Q(expired_count=0)) & Q(id__gt=last_id) + ).only( + "id", "code", "prefix", "suffix", "file_path", "uuid_file_name" + ).order_by("id").limit(CLEANUP_BATCH_SIZE) + if not expire_data: + break + for exp in expire_data: + last_id = exp.id + try: + if not is_local_ref(exp): + # 删除只需要路径字段；不加载 text 大字段 + await file_storage.delete_file( + StoredFile( + file_path=exp.file_path, + uuid_file_name=exp.uuid_file_name, + code=exp.code, + prefix=exp.prefix, + suffix=exp.suffix, + ) + ) + except Exception as e: + logger.error(f"删除过期文件失败 code={exp.code}: {e}") + try: + await exp.delete() + except Exception as e: + logger.error(f"删除记录失败 code={exp.code}: {e}") except Exception as e: logger.error(e) finally: @@ -71,33 +92,41 @@ async def clean_incomplete_uploads(): expire_hours = getattr(settings, "chunk_expire_hours", 24) now = await get_now() expire_time = now - datetime.timedelta(hours=expire_hours) - expired_sessions = await UploadChunk.filter( - chunk_index=-1, created_at__lt=expire_time - ).all() - - for session in expired_sessions: - try: - save_path = session.save_path - if not save_path: - _, _, _, _, save_path = await get_chunk_file_path_name( - session.file_name, session.upload_id + last_id = 0 + while True: + expired_sessions = await UploadChunk.filter( + chunk_index=-1, + created_at__lt=expire_time, + id__gt=last_id, + ).only("id", "upload_id", "save_path", "file_name").order_by( + "id" + ).limit(CLEANUP_BATCH_SIZE) + if not expired_sessions: + break + for session in expired_sessions: + last_id = session.id + try: + save_path = session.save_path + if not save_path: + _, _, _, _, save_path = await get_chunk_file_path_name( + session.file_name, session.upload_id + ) + await file_storage.clean_chunks(session.upload_id, save_path) + except Exception as e: + logger.error( + f"清理分片文件失败 upload_id={session.upload_id}: {e}" ) - await file_storage.clean_chunks(session.upload_id, save_path) - except Exception as e: - logger.error( - f"清理分片文件失败 upload_id={session.upload_id}: {e}" - ) - try: - await UploadChunk.filter(upload_id=session.upload_id).delete() - await StorageReservation.filter( - token=f"chunk:{session.upload_id}" - ).delete() - logger.info(f"已清理过期上传会话 upload_id={session.upload_id}") - except Exception as e: - logger.error( - f"删除分片记录失败 upload_id={session.upload_id}: {e}" - ) + try: + await UploadChunk.filter(upload_id=session.upload_id).delete() + await StorageReservation.filter( + token=f"chunk:{session.upload_id}" + ).delete() + logger.info(f"已清理过期上传会话 upload_id={session.upload_id}") + except Exception as e: + logger.error( + f"删除分片记录失败 upload_id={session.upload_id}: {e}" + ) except Exception as e: logger.error(f"清理未完成上传任务异常: {e}") @@ -111,28 +140,41 @@ async def clean_expired_presign_sessions(): await refresh_settings() storage: FileStorageInterface = storages[settings.file_storage]() now = await get_now() - expired_sessions = await PresignUploadSession.filter( - expires_at__lt=now - ).all() - for session in expired_sessions: - if session.mode == "direct": - try: - if await storage.file_exists(session.save_path): - await storage.delete_file( - StoredFile( - file_path=os.path.dirname(session.save_path), - uuid_file_name=os.path.basename(session.save_path), + last_id = 0 + while True: + expired_sessions = await PresignUploadSession.filter( + expires_at__lt=now, id__gt=last_id + ).only("id", "upload_id", "mode", "save_path").order_by( + "id" + ).limit(CLEANUP_BATCH_SIZE) + if not expired_sessions: + break + for session in expired_sessions: + last_id = session.id + if session.mode == "direct": + try: + if await storage.file_exists(session.save_path): + await storage.delete_file( + StoredFile( + file_path=os.path.dirname(session.save_path), + uuid_file_name=os.path.basename(session.save_path), + ) ) + except Exception as e: + logger.error( + "清理过期直传文件失败 " + f"upload_id={session.upload_id}: {e}" ) + try: + await StorageReservation.filter( + token=f"presign:{session.upload_id}" + ).delete() + await session.delete() except Exception as e: logger.error( - "清理过期直传文件失败 " + "删除过期预签名会话记录失败 " f"upload_id={session.upload_id}: {e}" ) - await StorageReservation.filter( - token=f"presign:{session.upload_id}" - ).delete() - await session.delete() except Exception as e: logger.error(f"清理过期预签名会话异常: {e}") finally: diff --git a/tests/test_cleanup_batching.py b/tests/test_cleanup_batching.py new file mode 100644 index 000000000..9264eaddc --- /dev/null +++ b/tests/test_cleanup_batching.py @@ -0,0 +1,122 @@ +"""过期清理任务游标分批行为。 + +验证超过单批上限的记录仍能在同一轮里全部处理，且单条删除失败不会让 +循环在同一批上反复重试（游标必须前进）。 +""" +import datetime +from unittest.mock import patch + +import pytest + +from apps.base.models import FileCodes, PresignUploadSession +from apps.base.tasks import clean_expired_presign_sessions, delete_expire_files +from core.utils import get_now + + +class _SleepSentinel(Exception): + pass + + +async def _run_round(tmp_path): + def _sleep(_seconds): + raise _SleepSentinel + + with patch("apps.base.tasks.data_root", tmp_path), patch( + "core.storage.local.data_root", tmp_path + ), patch("apps.base.tasks.asyncio.sleep", side_effect=_sleep): + with pytest.raises(_SleepSentinel): + await delete_expire_files() + + +async def _make_expired(code: str, tmp_path, now): + file_path = "share/data/2026/01/01" + path = tmp_path / file_path / f"uuid-{code}" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"payload") + return await FileCodes.create( + code=code, + file_path=file_path, + uuid_file_name=f"uuid-{code}", + size=1, + expired_at=now - datetime.timedelta(days=1), + ) + + +@pytest.mark.asyncio +async def test_more_than_one_batch_all_cleaned(db, tmp_path, monkeypatch): + monkeypatch.setattr("apps.base.tasks.CLEANUP_BATCH_SIZE", 2) + now = await get_now() + codes = [f"batch{i}" for i in range(5)] + for code in codes: + await _make_expired(code, tmp_path, now) + await FileCodes.create( + code="alive", + file_path="share/data/2026/01/01", + uuid_file_name="uuid-alive", + expired_at=now + datetime.timedelta(days=30), + ) + + await _run_round(tmp_path) + + assert await FileCodes.all().count() == 1 + assert (await FileCodes.all().first()).code == "alive" + + +@pytest.mark.asyncio +async def test_failed_row_does_not_block_progress(db, tmp_path, monkeypatch): + monkeypatch.setattr("apps.base.tasks.CLEANUP_BATCH_SIZE", 1) + now = await get_now() + for code in ["bad", "good1", "good2"]: + await _make_expired(code, tmp_path, now) + + real_delete = FileCodes.delete + + async def flaky_delete(self, using_db=None): + if self.code == "bad": + raise RuntimeError("delete boom") + return await real_delete(self, using_db=using_db) + + monkeypatch.setattr(FileCodes, "delete", flaky_delete) + + await _run_round(tmp_path) + + assert await FileCodes.filter(code="bad").exists() + assert not await FileCodes.filter(code="good1").exists() + assert not await FileCodes.filter(code="good2").exists() + + +@pytest.mark.asyncio +async def test_failed_presign_delete_does_not_block_progress(db, tmp_path, monkeypatch): + monkeypatch.setattr("apps.base.tasks.CLEANUP_BATCH_SIZE", 1) + now = await get_now() + for uid in ["bad-p", "good-p1", "good-p2"]: + await PresignUploadSession.create( + upload_id=uid, + file_name="f.bin", + file_size=1, + save_path="share/data/2026/01/01/f.bin", + mode="proxy", + expires_at=now - datetime.timedelta(minutes=5), + ) + + real_delete = PresignUploadSession.delete + + async def flaky_delete(self, using_db=None): + if self.upload_id == "bad-p": + raise RuntimeError("presign delete boom") + return await real_delete(self, using_db=using_db) + + monkeypatch.setattr(PresignUploadSession, "delete", flaky_delete) + + def _sleep(_seconds): + raise _SleepSentinel + + with patch("apps.base.tasks.data_root", tmp_path), patch( + "core.storage.local.data_root", tmp_path + ), patch("apps.base.tasks.asyncio.sleep", side_effect=_sleep): + with pytest.raises(_SleepSentinel): + await clean_expired_presign_sessions() + + assert await PresignUploadSession.filter(upload_id="bad-p").exists() + assert not await PresignUploadSession.filter(upload_id="good-p1").exists() + assert not await PresignUploadSession.filter(upload_id="good-p2").exists()