diff --git a/astrbot/core/backup/exporter.py b/astrbot/core/backup/exporter.py index a922375998..557315f181 100644 --- a/astrbot/core/backup/exporter.py +++ b/astrbot/core/backup/exporter.py @@ -4,10 +4,14 @@ 导出格式为 JSON，这是数据库无关的方案，支持未来向 MySQL/PostgreSQL 迁移。 """ +import asyncio +import functools import hashlib import json import os import zipfile +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING, Any @@ -64,6 +68,37 @@ def __init__( self.kb_manager = kb_manager self.config_path = config_path self._checksums: dict[str, str] = {} + self._archive_executor: ThreadPoolExecutor | None = None + + async def _run_in_archive_thread( + self, func: Callable[..., Any], /, *args: Any, **kwargs: Any + ) -> Any: + """Run an archive operation on the executor that owns the archive. + + Args: + func: The blocking archive callable to execute, e.g. ``zf.write``. + *args: Positional arguments forwarded to ``func``. + **kwargs: Keyword arguments forwarded to ``func``. + + Returns: + Whatever ``func`` returns. + + Raises: + RuntimeError: If the private archive helpers are called outside + ``export_all``, i.e. while no archive executor owns the archive. + """ + # The private archive helpers are only valid while export_all owns the + # executor. Fail loudly instead of silently falling back to run_in_executor(None), + # which would use the shared multi-threaded default pool and break serialisation. + if self._archive_executor is None: + raise RuntimeError( + "Internal misuse: _run_in_archive_thread was called outside " + "export_all, so no archive executor owns the archive." + ) + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._archive_executor, functools.partial(func, *args, **kwargs) + ) async def export_all( self, @@ -78,7 +113,26 @@ async def export_all( Returns: str: 生成的 ZIP 文件路径 + + Raises: + RuntimeError: If another export is already running on this instance. + BaseException: Whatever the export failed or was cancelled with. A failure to + finalise the archive (such as ``OSError`` from a full disk) is reported + here too, so a caller can never mistake an unopenable ZIP for a + completed backup. """ + # Re-entrancy guard, deliberately placed before any instance state is mutated. + # A second concurrent call would otherwise replace self._archive_executor while + # the first call is still using it, and whichever call finishes first destroys the + # executor the other one is still submitting to. Real call sites + # (dashboard/services/backup_service.py) build a fresh exporter per export, so + # refusing here costs nothing in production and keeps the failure mode obvious. + if self._archive_executor is not None: + raise RuntimeError( + "A backup export is already running on this exporter instance; create a " + "separate AstrBotExporter for each concurrent export." + ) + if output_dir is None: output_dir = get_astrbot_backups_path() @@ -91,117 +145,271 @@ async def export_all( logger.info(f"开始导出备份到 {zip_path}") + # Compression of multi-GB plugin data previously ran on the event loop and + # froze it for minutes. A single dedicated worker moves that work off the loop, + # and having exactly one worker serialises every access to the ZipFile, which + # is not safe for concurrent use. + archive_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="astrbot-backup-zip" + ) + self._archive_executor = archive_executor + + zf: zipfile.ZipFile | None = None + open_task: asyncio.Future | None = None + # Set only when the export itself raises. Kept separate from the close result so + # the teardown can tell "the archive was never finished" apart from "the archive + # finished but could not be finalised". + export_error: BaseException | None = None + close_error: BaseException | None = None + cancelled_while_cleaning = False + try: - with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: - # 1. 导出主数据库 + # The ZipFile is created on the archive worker instead of on the event loop. + # The matching close happens on that same worker in the `finally` below, and + # because the executor has exactly one worker (run_in_executor is a FIFO + # queue served by it), a close submitted after an in-flight write is + # guaranteed to run after that write finishes and can never overlap it. + # Closing from the event loop instead would call ZipFile.close() while a + # zf.write is still running, which raises "Can't close the ZIP file while + # there is an open writing handle on it" and leaves an unusable partial zip. + open_task = asyncio.ensure_future( + self._run_in_archive_thread( + zipfile.ZipFile, zip_path, "w", zipfile.ZIP_DEFLATED + ) + ) + zf = await asyncio.shield(open_task) + + # 1. 导出主数据库 + if progress_callback: + await progress_callback("main_db", 0, 100, "正在导出主数据库...") + main_data = await self._export_main_database() + main_db_json = json.dumps( + main_data, ensure_ascii=False, indent=2, default=str + ) + await self._run_in_archive_thread( + zf.writestr, "databases/main_db.json", main_db_json + ) + self._add_checksum("databases/main_db.json", main_db_json) + if progress_callback: + await progress_callback("main_db", 100, 100, "主数据库导出完成") + + # 2. 导出知识库数据 + kb_meta_data: dict[str, Any] = { + "knowledge_bases": [], + "kb_documents": [], + "kb_media": [], + } + if self.kb_manager: if progress_callback: - await progress_callback("main_db", 0, 100, "正在导出主数据库...") - main_data = await self._export_main_database() - main_db_json = json.dumps( - main_data, ensure_ascii=False, indent=2, default=str + await progress_callback( + "kb_metadata", 0, 100, "正在导出知识库元数据..." + ) + kb_meta_data = await self._export_kb_metadata() + kb_meta_json = json.dumps( + kb_meta_data, ensure_ascii=False, indent=2, default=str ) - zf.writestr("databases/main_db.json", main_db_json) - self._add_checksum("databases/main_db.json", main_db_json) + await self._run_in_archive_thread( + zf.writestr, "databases/kb_metadata.json", kb_meta_json + ) + self._add_checksum("databases/kb_metadata.json", kb_meta_json) if progress_callback: - await progress_callback("main_db", 100, 100, "主数据库导出完成") - - # 2. 导出知识库数据 - kb_meta_data: dict[str, Any] = { - "knowledge_bases": [], - "kb_documents": [], - "kb_media": [], - } - if self.kb_manager: - if progress_callback: - await progress_callback( - "kb_metadata", 0, 100, "正在导出知识库元数据..." - ) - kb_meta_data = await self._export_kb_metadata() - kb_meta_json = json.dumps( - kb_meta_data, ensure_ascii=False, indent=2, default=str + await progress_callback( + "kb_metadata", 100, 100, "知识库元数据导出完成" ) - zf.writestr("databases/kb_metadata.json", kb_meta_json) - self._add_checksum("databases/kb_metadata.json", kb_meta_json) - if progress_callback: - await progress_callback( - "kb_metadata", 100, 100, "知识库元数据导出完成" - ) - - # 导出每个知识库的文档数据 - kb_insts = self.kb_manager.kb_insts - total_kbs = len(kb_insts) - for idx, (kb_id, kb_helper) in enumerate(kb_insts.items()): - if progress_callback: - await progress_callback( - "kb_documents", - idx, - total_kbs, - f"正在导出知识库 {kb_helper.kb.kb_name} 的文档数据...", - ) - doc_data = await self._export_kb_documents(kb_helper) - doc_json = json.dumps( - doc_data, ensure_ascii=False, indent=2, default=str - ) - doc_path = f"databases/kb_{kb_id}/documents.json" - zf.writestr(doc_path, doc_json) - self._add_checksum(doc_path, doc_json) - - # 导出 FAISS 索引文件 - await self._export_faiss_index(zf, kb_helper, kb_id) - - # 导出知识库多媒体文件 - await self._export_kb_media_files(zf, kb_helper, kb_id) + # 导出每个知识库的文档数据 + kb_insts = self.kb_manager.kb_insts + total_kbs = len(kb_insts) + for idx, (kb_id, kb_helper) in enumerate(kb_insts.items()): if progress_callback: await progress_callback( - "kb_documents", total_kbs, total_kbs, "知识库文档导出完成" + "kb_documents", + idx, + total_kbs, + f"正在导出知识库 {kb_helper.kb.kb_name} 的文档数据...", ) + doc_data = await self._export_kb_documents(kb_helper) + doc_json = json.dumps( + doc_data, ensure_ascii=False, indent=2, default=str + ) + doc_path = f"databases/kb_{kb_id}/documents.json" + await self._run_in_archive_thread(zf.writestr, doc_path, doc_json) + self._add_checksum(doc_path, doc_json) - # 3. 导出配置文件 - if progress_callback: - await progress_callback("config", 0, 100, "正在导出配置文件...") - if os.path.exists(self.config_path): - with open(self.config_path, encoding="utf-8") as f: - config_content = f.read() - zf.writestr("config/cmd_config.json", config_content) - self._add_checksum("config/cmd_config.json", config_content) - if progress_callback: - await progress_callback("config", 100, 100, "配置文件导出完成") + # 导出 FAISS 索引文件 + await self._export_faiss_index(zf, kb_helper, kb_id) - # 4. 导出附件文件 - if progress_callback: - await progress_callback("attachments", 0, 100, "正在导出附件...") - await self._export_attachments(zf, main_data.get("attachments", [])) - if progress_callback: - await progress_callback("attachments", 100, 100, "附件导出完成") + # 导出知识库多媒体文件 + await self._export_kb_media_files(zf, kb_helper, kb_id) - # 5. 导出插件和其他目录 if progress_callback: await progress_callback( - "directories", 0, 100, "正在导出插件和数据目录..." + "kb_documents", total_kbs, total_kbs, "知识库文档导出完成" ) - dir_stats = await self._export_directories(zf) - if progress_callback: - await progress_callback("directories", 100, 100, "目录导出完成") - # 6. 生成 manifest - if progress_callback: - await progress_callback("manifest", 0, 100, "正在生成清单...") - manifest = self._generate_manifest(main_data, kb_meta_data, dir_stats) - manifest_json = json.dumps(manifest, ensure_ascii=False, indent=2) - zf.writestr("manifest.json", manifest_json) - if progress_callback: - await progress_callback("manifest", 100, 100, "清单生成完成") - - logger.info(f"备份导出完成: {zip_path}") - return zip_path + # 3. 导出配置文件 + if progress_callback: + await progress_callback("config", 0, 100, "正在导出配置文件...") + if os.path.exists(self.config_path): + with open(self.config_path, encoding="utf-8") as f: + config_content = f.read() + await self._run_in_archive_thread( + zf.writestr, "config/cmd_config.json", config_content + ) + self._add_checksum("config/cmd_config.json", config_content) + if progress_callback: + await progress_callback("config", 100, 100, "配置文件导出完成") + + # 4. 导出附件文件 + if progress_callback: + await progress_callback("attachments", 0, 100, "正在导出附件...") + await self._export_attachments(zf, main_data.get("attachments", [])) + if progress_callback: + await progress_callback("attachments", 100, 100, "附件导出完成") + + # 5. 导出插件和其他目录 + if progress_callback: + await progress_callback( + "directories", 0, 100, "正在导出插件和数据目录..." + ) + dir_stats = await self._export_directories(zf) + if progress_callback: + await progress_callback("directories", 100, 100, "目录导出完成") + + # 6. 生成 manifest + if progress_callback: + await progress_callback("manifest", 0, 100, "正在生成清单...") + manifest = self._generate_manifest(main_data, kb_meta_data, dir_stats) + manifest_json = json.dumps(manifest, ensure_ascii=False, indent=2) + await self._run_in_archive_thread( + zf.writestr, "manifest.json", manifest_json + ) + if progress_callback: + await progress_callback("manifest", 100, 100, "清单生成完成") except Exception as e: + # Only real failures are logged here. asyncio.CancelledError does not inherit + # from Exception, so a cancelled export falls straight through to the cleanup + # below instead of being reported as a failed backup. logger.error(f"备份导出失败: {e}") - # 清理失败的文件 - if os.path.exists(zip_path): - os.remove(zip_path) + export_error = e raise + except BaseException as e: + # A cancellation that arrives while the export is still producing data. It is + # remembered so the teardown can tell it apart from an export that produced a + # finished archive, and re-raised untouched. + export_error = e + raise + finally: + # Teardown never returns early. A cancellation delivered while the archive is + # still being closed is remembered here and re-raised only after every + # cleanup step has run, so a partial zip is never left behind and the worker + # is never left running with self._archive_executor still set. + try: + # Waiting on a shield is the whole cancellation story here. asyncio.shield + # raises the caller's CancelledError immediately while the inner operation + # keeps running on the archive worker, so the wait continues until that + # operation truly finishes; abandoning it would leave the archive half + # closed. Nothing in this class ever cancels these two tasks, so any + # CancelledError raised at these awaits is this caller's own cancellation. + if open_task is not None: + while not open_task.done(): + try: + await asyncio.shield(open_task) + except asyncio.CancelledError: + cancelled_while_cleaning = True + except Exception: + # The open itself failed. Its result is read below to report + # why, so swallow it here and let the teardown continue. + pass + if not open_task.cancelled(): + try: + # The archive exists even though zf was never assigned when the + # export was cancelled while ZipFile() was still opening. Take + # the instance so its handle gets closed instead of leaked. + if zf is None: + zf = open_task.result() + except Exception as e: + logger.warning(f"打开备份归档失败: {e}") + + if zf is not None: + # Close the archive on its own worker. The close is submitted after the + # last write on the same FIFO single-worker executor, so it can never + # overlap a write that is still in flight, and awaiting it keeps the + # event loop responsive while a large compression finishes. A worker + # that hangs forever is not bounded here; the caller waits for it so + # the executor can always be shut down without a dangling writer. + close_task = asyncio.ensure_future( + self._run_in_archive_thread(zf.close) + ) + while not close_task.done(): + try: + await asyncio.shield(close_task) + except asyncio.CancelledError: + cancelled_while_cleaning = True + except Exception: + # A failed close is an export failure, reported below. + pass + if not close_task.cancelled(): + try: + close_task.result() + except BaseException as e: + close_error = e + finally: + # Everything else in the teardown is best-effort, but this is not: a + # BaseException escaping the waits above (KeyboardInterrupt, SystemExit) + # must still release the instance, or every later export_all on it would + # be refused as "already running" forever. The worker has drained the + # archive by now, so shutdown(wait=True) joins an idle thread and returns + # promptly instead of stalling the event loop behind a compression. + self._archive_executor = None + archive_executor.shutdown(wait=True) + + if ( + export_error is not None + or close_error is not None + or cancelled_while_cleaning + ): + # This call is not going to report a usable archive: the export failed, it + # was cancelled, or the archive could not be finalised and is left without + # its central directory. Remove the file so no caller can pick up a + # truncated zip that looks usable, and so a cancelled export does not + # orphan a file nothing will ever reference. Removal is best-effort and + # must never mask the error raised below; ZipFile.close() always clears + # its handle in a `finally`, so the usual reason a removal fails here is + # an unrelated filesystem condition. + if os.path.exists(zip_path): + try: + os.remove(zip_path) + except OSError as remove_error: + logger.warning(f"清理残缺备份文件失败: {remove_error}") + + if close_error is not None: + # The archive could not be finalised, so it is unopenable. Reporting + # success here would tell the caller a corrupt backup was written, which + # is the worst outcome for a backup feature: this is an export failure. + logger.error(f"备份归档收尾失败，导出按失败处理: {close_error}") + if export_error is not None: + # Keep whatever was already in flight reachable as the cause. + raise close_error from export_error + raise close_error + + if export_error is None: + # Nothing else is propagating, so the teardown's own bookkeeping decides + # the outcome. Otherwise the exception from the try block resumes by + # itself, with its original traceback (and, for a cancellation, with the + # task still reported as cancelled). + if cancelled_while_cleaning: + # Restore the cancellation that the teardown deliberately absorbed, so + # the caller still observes this export as cancelled. It is + # deliberately not restored when the export already failed: that error + # is the more useful report, and a bare CancelledError would hide it + # from the dashboard's `except Exception`, leaving the task stuck in + # "processing". + raise asyncio.CancelledError + logger.info(f"备份导出完成: {zip_path}") + + return zip_path async def _export_main_database(self) -> dict[str, list[dict]]: """导出主数据库所有表""" @@ -280,7 +488,9 @@ async def _export_faiss_index( index_path = kb_helper.kb_dir / "index.faiss" if index_path.exists(): archive_path = f"databases/kb_{kb_id}/index.faiss" - zf.write(str(index_path), archive_path) + await self._run_in_archive_thread( + zf.write, str(index_path), archive_path + ) logger.debug(f"导出 FAISS 索引: {archive_path}") except Exception as e: logger.warning(f"导出 FAISS 索引失败: {e}") @@ -289,87 +499,26 @@ async def _export_kb_media_files( self, zf: zipfile.ZipFile, kb_helper: Any, kb_id: str ) -> None: """导出知识库的多媒体文件""" - try: - media_dir = kb_helper.kb_medias_dir - if not media_dir.exists(): - return - - for root, _, files in os.walk(media_dir): - for file in files: - file_path = Path(root) / file - # 计算相对路径 - rel_path = file_path.relative_to(kb_helper.kb_dir) - archive_path = f"files/kb_media/{kb_id}/{rel_path}" - zf.write(str(file_path), archive_path) - except Exception as e: - logger.warning(f"导出知识库媒体文件失败: {e}") + await self._run_in_archive_thread(_write_kb_media_files, zf, kb_helper, kb_id) async def _export_directories( self, zf: zipfile.ZipFile ) -> dict[str, dict[str, int]]: """导出插件和其他数据目录 + Args: + zf: 备份 ZIP 归档对象，所有写入都在归档线程中串行执行 + Returns: dict: 每个目录的统计信息 {dir_name: {"files": count, "size": bytes}} """ - stats: dict[str, dict[str, int]] = {} - backup_directories = get_backup_directories() - - for dir_name, dir_path in backup_directories.items(): - full_path = Path(dir_path) - if not full_path.exists(): - logger.debug(f"目录不存在，跳过: {full_path}") - continue - - file_count = 0 - total_size = 0 - - try: - for root, dirs, files in os.walk(full_path): - # 跳过 __pycache__ 目录 - dirs[:] = [d for d in dirs if d != "__pycache__"] - - for file in files: - # 跳过 .pyc 文件 - if file.endswith(".pyc"): - continue - - file_path = Path(root) / file - try: - # 计算相对路径 - rel_path = file_path.relative_to(full_path) - archive_path = f"directories/{dir_name}/{rel_path}" - zf.write(str(file_path), archive_path) - file_count += 1 - total_size += file_path.stat().st_size - except Exception as e: - logger.warning(f"导出文件 {file_path} 失败: {e}") - - stats[dir_name] = {"files": file_count, "size": total_size} - logger.debug( - f"导出目录 {dir_name}: {file_count} 个文件, {total_size} 字节" - ) - except Exception as e: - logger.warning(f"导出目录 {dir_path} 失败: {e}") - stats[dir_name] = {"files": 0, "size": 0} - - return stats + return await self._run_in_archive_thread(_write_backup_directories, zf) async def _export_attachments( self, zf: zipfile.ZipFile, attachments: list[dict] ) -> None: """导出附件文件""" - for attachment in attachments: - try: - file_path = attachment.get("path", "") - if file_path and os.path.exists(file_path): - # 使用 attachment_id 作为文件名 - attachment_id = attachment.get("attachment_id", "") - ext = os.path.splitext(file_path)[1] - archive_path = f"files/attachments/{attachment_id}{ext}" - zf.write(file_path, archive_path) - except Exception as e: - logger.warning(f"导出附件失败: {e}") + await self._run_in_archive_thread(_write_attachments, zf, attachments) def _model_to_dict(self, record: Any) -> dict: """将 SQLModel 实例转换为字典 @@ -475,3 +624,101 @@ def _generate_manifest( } return manifest + + +def _write_backup_directories(zf: zipfile.ZipFile) -> dict[str, dict[str, int]]: + """将插件和其他数据目录写入备份归档。 + + 这是阻塞式目录遍历与压缩逻辑，必须通过 ``AstrBotExporter._run_in_archive_thread`` + 在归档线程中执行，避免冻结事件循环。 + + Args: + zf: 备份 ZIP 归档对象，写入操作由调用方保证串行执行 + + Returns: + dict: 每个目录的统计信息 {dir_name: {"files": count, "size": bytes}} + """ + stats: dict[str, dict[str, int]] = {} + backup_directories = get_backup_directories() + + for dir_name, dir_path in backup_directories.items(): + full_path = Path(dir_path) + if not full_path.exists(): + logger.debug(f"目录不存在，跳过: {full_path}") + continue + + file_count = 0 + total_size = 0 + + try: + for root, dirs, files in os.walk(full_path): + # 跳过 __pycache__ 目录 + dirs[:] = [d for d in dirs if d != "__pycache__"] + + for file in files: + # 跳过 .pyc 文件 + if file.endswith(".pyc"): + continue + + file_path = Path(root) / file + try: + # 计算相对路径 + rel_path = file_path.relative_to(full_path) + archive_path = f"directories/{dir_name}/{rel_path}" + zf.write(str(file_path), archive_path) + file_count += 1 + total_size += file_path.stat().st_size + except Exception as e: + logger.warning(f"导出文件 {file_path} 失败: {e}") + + stats[dir_name] = {"files": file_count, "size": total_size} + logger.debug(f"导出目录 {dir_name}: {file_count} 个文件, {total_size} 字节") + except Exception as e: + logger.warning(f"导出目录 {dir_path} 失败: {e}") + stats[dir_name] = {"files": 0, "size": 0} + + return stats + + +def _write_attachments(zf: zipfile.ZipFile, attachments: list[dict]) -> None: + """将附件文件写入备份归档。 + + Args: + zf: 备份 ZIP 归档对象，写入操作由调用方保证串行执行 + attachments: 附件记录列表，每项包含 path 与 attachment_id 字段 + """ + for attachment in attachments: + try: + file_path = attachment.get("path", "") + if file_path and os.path.exists(file_path): + # 使用 attachment_id 作为文件名 + attachment_id = attachment.get("attachment_id", "") + ext = os.path.splitext(file_path)[1] + archive_path = f"files/attachments/{attachment_id}{ext}" + zf.write(file_path, archive_path) + except Exception as e: + logger.warning(f"导出附件失败: {e}") + + +def _write_kb_media_files(zf: zipfile.ZipFile, kb_helper: Any, kb_id: str) -> None: + """将知识库多媒体文件写入备份归档。 + + Args: + zf: 备份 ZIP 归档对象，写入操作由调用方保证串行执行 + kb_helper: 知识库实例，提供 kb_dir 与 kb_medias_dir 路径 + kb_id: 知识库 ID，用于生成归档内的相对路径 + """ + try: + media_dir = kb_helper.kb_medias_dir + if not media_dir.exists(): + return + + for root, _, files in os.walk(media_dir): + for file in files: + file_path = Path(root) / file + # 计算相对路径 + rel_path = file_path.relative_to(kb_helper.kb_dir) + archive_path = f"files/kb_media/{kb_id}/{rel_path}" + zf.write(str(file_path), archive_path) + except Exception as e: + logger.warning(f"导出知识库媒体文件失败: {e}") diff --git a/tests/test_backup.py b/tests/test_backup.py index 069b61d6d6..9d91fbefef 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -1,8 +1,13 @@ """备份功能单元测试""" +import asyncio +import errno +import hashlib import json import os import re +import threading +import time import zipfile from datetime import datetime from pathlib import Path @@ -1458,3 +1463,663 @@ async def test_short_write_chunk_is_rejected(self, backup_service): {"upload_id": upload_id}, owner="alice" ) assert result["size"] == 100 + + +class TestBackupExportDoesNotBlockEventLoop: + """测试备份导出不会阻塞事件循环（归档操作在专用工作线程中执行）""" + + @staticmethod + def _make_exporter(data_dir: Path) -> AstrBotExporter: + """构造一个不访问真实用户数据的导出器 + + Args: + data_dir: 作为备份数据源的临时数据目录 + + Returns: + AstrBotExporter: 只依赖 mock 数据库与临时配置目录的导出器实例 + """ + session = AsyncMock() + result = MagicMock() + result.scalars.return_value.all.return_value = [] + session.execute = AsyncMock(return_value=result) + + mock_db = MagicMock() + mock_db.get_db.return_value = AsyncMock( + __aenter__=AsyncMock(return_value=session), + __aexit__=AsyncMock(return_value=None), + ) + return AstrBotExporter( + main_db=mock_db, + kb_manager=None, + config_path=str(data_dir / "cmd_config.json"), + ) + + @staticmethod + def _make_compressible_plugin_dir(data_dir: Path, blob_count: int = 4) -> Path: + """创建不可压缩的大文件目录，保证 DEFLATE 有可测量的计算量 + + Args: + data_dir: 作为备份数据源的临时数据目录 + blob_count: 生成的随机数据文件数量，每个 16 MiB + + Returns: + Path: 插件目录路径 + """ + plugin_dir = data_dir / "plugins" + plugin_dir.mkdir(exist_ok=True) + for index in range(blob_count): + (plugin_dir / f"blob_{index}.bin").write_bytes(os.urandom(16 * 1024 * 1024)) + return plugin_dir + + @staticmethod + async def _watch_event_loop_gaps(gaps: list[float], stop: asyncio.Event) -> None: + """以 5ms 心跳采样事件循环停顿，写入 ``gaps`` + + Args: + gaps: 每次心跳实际间隔的收集列表 + stop: 置位后心跳退出 + """ + last = time.monotonic() + while not stop.is_set(): + await asyncio.sleep(0.005) + now = time.monotonic() + gaps.append(now - last) + last = now + + @pytest.mark.asyncio + async def test_archive_writes_leave_event_loop_thread(self, tmp_path, monkeypatch): + """归档写入必须发生在非事件循环线程中（线程守卫回归测试）""" + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = data_dir / "plugins" + plugin_dir.mkdir() + for index in range(5): + (plugin_dir / f"plugin_{index}.txt").write_text(f"content-{index}" * 100) + (plugin_dir / "nested").mkdir() + (plugin_dir / "nested" / "deep.txt").write_text("deep content" * 100) + + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(plugin_dir)}, + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + write_threads: set[int] = set() + real_write = zipfile.ZipFile.write + real_writestr = zipfile.ZipFile.writestr + + def spy_write(*args, **kwargs): + write_threads.add(threading.get_ident()) + return real_write(*args, **kwargs) + + def spy_writestr(*args, **kwargs): + write_threads.add(threading.get_ident()) + return real_writestr(*args, **kwargs) + + monkeypatch.setattr(zipfile.ZipFile, "write", spy_write) + monkeypatch.setattr(zipfile.ZipFile, "writestr", spy_writestr) + + loop_thread_id = threading.get_ident() + zip_path = await exporter.export_all(output_dir=str(backup_dir)) + + assert os.path.exists(zip_path) + # 归档写入确实发生过，且全部离开了事件循环线程 + assert write_threads + assert loop_thread_id not in write_threads + + # 产物仍然完整可读 + with zipfile.ZipFile(zip_path, "r") as zf: + namelist = zf.namelist() + assert "manifest.json" in namelist + assert "directories/plugins/plugin_0.txt" in namelist + assert "directories/plugins/nested/deep.txt" in namelist + + @pytest.mark.asyncio + async def test_event_loop_stays_responsive_during_export( + self, tmp_path, monkeypatch + ): + """导出压缩期间事件循环心跳既要有足够采样、也不能长时间停顿 + + 采样数量断言是回归守卫的核心：未修复的实现会把事件循环整段冻结，心跳任务 + 一次都跑不到，只在导出结束后才采样到一两段间隔。没有这条断言时，空采样集 + 会让下面的停顿断言退化成恒真。 + """ + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = self._make_compressible_plugin_dir(data_dir) + + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(plugin_dir)}, + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + gaps: list[float] = [] + stop_heartbeat = asyncio.Event() + heartbeat_task = asyncio.create_task( + self._watch_event_loop_gaps(gaps, stop_heartbeat) + ) + + started = time.monotonic() + try: + zip_path = await exporter.export_all(output_dir=str(backup_dir)) + finally: + # 采样窗口包含导出结束后的收尾阶段，因此覆盖的是整段导出墙钟时间 + wall_time = time.monotonic() - started + stop_heartbeat.set() + await heartbeat_task + + assert os.path.exists(zip_path) + + # 事件循环必须真的被调度过：阻塞式实现下心跳一次都跑不到 + assert len(gaps) >= 20, f"事件循环仅采样到 {len(gaps)} 次心跳，导出期间被冻结" + covered = sum(gaps) + assert covered > wall_time * 0.8, ( + f"心跳仅覆盖 {covered:.3f}s / 导出墙钟 {wall_time:.3f}s" + ) + + max_gap = max(gaps) + assert max_gap < 0.5, f"事件循环最长停顿 {max_gap:.3f}s" + + @pytest.mark.asyncio + async def test_export_failure_shuts_down_executor_and_cleans_zip( + self, tmp_path, monkeypatch + ): + """导出失败时异常向上抛出、执行器被回收且不留残缺 ZIP""" + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = data_dir / "plugins" + plugin_dir.mkdir() + (plugin_dir / "plugin.txt").write_text("content") + + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(plugin_dir)}, + ) + + def boom(zf): + raise RuntimeError("simulated directory export failure") + + monkeypatch.setattr( + "astrbot.core.backup.exporter._write_backup_directories", boom + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + with pytest.raises(RuntimeError, match="simulated directory export failure"): + await exporter.export_all(output_dir=str(backup_dir)) + + # 执行器已释放，且失败路径清理了残缺 ZIP + assert exporter._archive_executor is None + assert list(backup_dir.glob("*.zip")) == [] + + @pytest.mark.asyncio + async def test_archive_helper_requires_active_executor(self): + """未处于 export_all 中时归档转发必须显式报错，而非回退到默认线程池""" + exporter = AstrBotExporter(main_db=MagicMock()) + + with pytest.raises(RuntimeError, match="export_all"): + await exporter._run_in_archive_thread(lambda: None) + + @pytest.mark.asyncio + async def test_concurrent_export_on_same_instance_is_rejected( + self, tmp_path, monkeypatch + ): + """同一实例上的并发导出必须立即失败，且不影响首个导出""" + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = self._make_compressible_plugin_dir(data_dir, blob_count=1) + + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(plugin_dir)}, + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + # 让首个导出停在归档线程里，保证第二个调用确实与之并发 + write_started = threading.Event() + release_write = threading.Event() + real_write = zipfile.ZipFile.write + + def blocking_write(*args, **kwargs): + write_started.set() + release_write.wait(timeout=30) + return real_write(*args, **kwargs) + + monkeypatch.setattr(zipfile.ZipFile, "write", blocking_write) + + first_task = asyncio.create_task( + exporter.export_all(output_dir=str(backup_dir)) + ) + try: + assert await asyncio.to_thread(write_started.wait, 30) + + with pytest.raises(RuntimeError, match="already running"): + await exporter.export_all(output_dir=str(backup_dir)) + + # 拒绝发生在动用实例状态之前：首个导出的执行器仍然有效 + assert exporter._archive_executor is not None + finally: + release_write.set() + + zip_path = await first_task + assert os.path.exists(zip_path) + assert exporter._archive_executor is None + + # 只有首个导出产物，被拒绝的调用没有留下残缺 ZIP + assert [path.name for path in backup_dir.glob("*.zip")] == [ + os.path.basename(zip_path) + ] + with zipfile.ZipFile(zip_path) as zf: + assert zf.testzip() is None + + @pytest.mark.asyncio + async def test_cancellation_closes_archive_and_keeps_loop_responsive( + self, tmp_path, monkeypatch + ): + """导出中途取消：归档被正确关闭、不留残缺 ZIP、事件循环不被拖住""" + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = self._make_compressible_plugin_dir(data_dir, blob_count=1) + + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(plugin_dir)}, + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + write_started = asyncio.Event() + real_write = zipfile.ZipFile.write + + def signalling_write(*args, **kwargs): + write_started.set() + return real_write(*args, **kwargs) + + monkeypatch.setattr(zipfile.ZipFile, "write", signalling_write) + + gaps: list[float] = [] + stop_heartbeat = asyncio.Event() + heartbeat_task = asyncio.create_task( + self._watch_event_loop_gaps(gaps, stop_heartbeat) + ) + + export_task = asyncio.create_task( + exporter.export_all(output_dir=str(backup_dir)) + ) + try: + await asyncio.wait_for(write_started.wait(), timeout=30) + gaps.clear() + + export_task.cancel() + with pytest.raises(asyncio.CancelledError): + await export_task + cancelled_gaps = list(gaps) + finally: + stop_heartbeat.set() + await heartbeat_task + + # 取消后的收尾不能把事件循环拖住（旧实现会同步等待压缩结束） + assert cancelled_gaps + assert max(cancelled_gaps) < 0.5, ( + f"取消期间事件循环最长停顿 {max(cancelled_gaps):.3f}s" + ) + + # 残缺 ZIP 已被清理，不会留下看似可读的损坏文件 + assert list(backup_dir.glob("*.zip")) == [] + assert exporter._archive_executor is None + assert not [ + thread + for thread in threading.enumerate() + if thread.name.startswith("astrbot-backup-zip") + ] + + # 同一数据目录上的新实例仍能正常导出 + retry_exporter = self._make_exporter(data_dir) + zip_path = await retry_exporter.export_all(output_dir=str(backup_dir)) + with zipfile.ZipFile(zip_path) as zf: + assert zf.testzip() is None + assert "manifest.json" in zf.namelist() + + @pytest.mark.asyncio + async def test_export_structure_is_stable(self, tmp_path, monkeypatch): + """归档结构稳定：成员清单一致、校验和与内容匹配、目录统计正确""" + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = data_dir / "plugins" + plugin_dir.mkdir() + (plugin_dir / "plugin.txt").write_text("content" * 10) + (plugin_dir / "nested").mkdir() + (plugin_dir / "nested" / "deep.txt").write_text("deep" * 10) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(plugin_dir)}, + ) + + first_dir = tmp_path / "backups-a" + second_dir = tmp_path / "backups-b" + first_zip = await self._make_exporter(data_dir).export_all( + output_dir=str(first_dir) + ) + second_zip = await self._make_exporter(data_dir).export_all( + output_dir=str(second_dir) + ) + + with ( + zipfile.ZipFile(first_zip) as first, + zipfile.ZipFile(second_zip) as second, + ): + # 归档完整可用，且同一份数据的两次导出结构完全一致 + assert first.testzip() is None + assert second.testzip() is None + assert first.namelist() == second.namelist() + assert set(first.namelist()) == { + "databases/main_db.json", + "config/cmd_config.json", + "directories/plugins/plugin.txt", + "directories/plugins/nested/deep.txt", + "manifest.json", + } + + first_manifest = json.loads(first.read("manifest.json")) + second_manifest = json.loads(second.read("manifest.json")) + # exported_at 每次不同，其余结构必须逐字节一致 + first_manifest.pop("exported_at") + second_manifest.pop("exported_at") + assert first_manifest == second_manifest + + # manifest 中的校验和与归档内实际内容逐项一致 + assert first_manifest["checksums"] + for path, checksum in first_manifest["checksums"].items(): + assert path in first.namelist() + digest = hashlib.sha256(first.read(path)).hexdigest() + assert checksum == f"sha256:{digest}" + + assert first_manifest["statistics"]["directories"] == { + "plugins": {"files": 2, "size": 110} + } + + +class _TeardownInterrupt(BaseException): + """模拟 KeyboardInterrupt/SystemExit 这类不继承 Exception 的收尾中断""" + + +class TestBackupExportTeardown: + """导出收尾语义：归档收尾失败与取消都不得掩盖真正的导出错误""" + + @staticmethod + def _make_exporter(data_dir: Path) -> AstrBotExporter: + """构造一个不访问真实用户数据的导出器 + + Args: + data_dir: 作为备份数据源的临时数据目录 + + Returns: + AstrBotExporter: 只依赖 mock 数据库与临时配置目录的导出器实例 + """ + session = AsyncMock() + result = MagicMock() + result.scalars.return_value.all.return_value = [] + session.execute = AsyncMock(return_value=result) + + mock_db = MagicMock() + mock_db.get_db.return_value = AsyncMock( + __aenter__=AsyncMock(return_value=session), + __aexit__=AsyncMock(return_value=None), + ) + return AstrBotExporter( + main_db=mock_db, + kb_manager=None, + config_path=str(data_dir / "cmd_config.json"), + ) + + @staticmethod + def _make_data_dir(tmp_path: Path) -> Path: + """创建带配置文件与插件目录的临时数据目录 + + Args: + tmp_path: pytest 提供的临时目录 + + Returns: + Path: 已填充的数据目录 + """ + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "cmd_config.json").write_text(json.dumps({"test": "config"})) + plugin_dir = data_dir / "plugins" + plugin_dir.mkdir() + (plugin_dir / "plugin.txt").write_text("content") + return data_dir + + @staticmethod + def _fail_end_record(*args, **kwargs): + """模拟中央目录写入失败，即 ENOSPC/EIO/配额耗尽真正暴露的位置""" + raise OSError(errno.ENOSPC, "No space left on device") + + @pytest.mark.asyncio + async def test_close_failure_is_reported_as_export_failure( + self, tmp_path, monkeypatch + ): + """收尾失败时 export_all 必须抛异常、清理残缺 ZIP 并释放执行器 + + 修复前该失败只被记成日志：export_all 正常返回路径，磁盘上留下一个 + 无法打开的 ZIP，调用方据此上报“备份成功”，用户在导入时才发现损坏。 + """ + data_dir = self._make_data_dir(tmp_path) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(data_dir / "plugins")}, + ) + monkeypatch.setattr(zipfile.ZipFile, "_write_end_record", self._fail_end_record) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + with pytest.raises(OSError) as excinfo: + await exporter.export_all(output_dir=str(backup_dir)) + + assert excinfo.value.errno == errno.ENOSPC + # 残缺 ZIP 不得留在输出目录，否则调用方会拿它当成功产物上报 + assert list(backup_dir.glob("*.zip")) == [] + # 实例必须仍然可用，不能因为收尾失败而永远卡在“already running” + # 即：正常结束时执行器已被释放，再入保护不会拒绝后续调用 + with monkeypatch.context() as retry_finalisation: + retry_finalisation.setattr( + zipfile.ZipFile, "_write_end_record", self._fail_end_record + ) + with pytest.raises(OSError) as retry_excinfo: + await exporter.export_all(output_dir=str(backup_dir)) + assert retry_excinfo.value.errno == errno.ENOSPC + assert list(backup_dir.glob("*.zip")) == [] + + @pytest.mark.asyncio + async def test_close_failure_keeps_original_error_as_cause( + self, tmp_path, monkeypatch + ): + """导出本身出错且收尾也失败时，两个错误都要能追溯""" + data_dir = self._make_data_dir(tmp_path) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(data_dir / "plugins")}, + ) + + def boom(zf): + raise RuntimeError("simulated directory export failure") + + monkeypatch.setattr( + "astrbot.core.backup.exporter._write_backup_directories", boom + ) + monkeypatch.setattr(zipfile.ZipFile, "_write_end_record", self._fail_end_record) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + with pytest.raises(OSError) as excinfo: + await exporter.export_all(output_dir=str(backup_dir)) + + assert excinfo.value.errno == errno.ENOSPC + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "simulated directory export failure" in str(excinfo.value.__cause__) + assert list(backup_dir.glob("*.zip")) == [] + + @pytest.mark.asyncio + async def test_base_exception_in_teardown_still_releases_instance( + self, tmp_path, monkeypatch + ): + """收尾抛出 BaseException 时执行器仍必须释放，实例不能被永久锁死 + + 修复前收尾只捕获 Exception：BaseException（KeyboardInterrupt/SystemExit） + 会直接冲出 finally，跳过 ``self._archive_executor = None``，此后同一实例上的 + 每次 export_all 都会以“already running”被拒绝。 + """ + data_dir = self._make_data_dir(tmp_path) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(data_dir / "plugins")}, + ) + + def interrupt_end_record(*args, **kwargs): + raise _TeardownInterrupt("simulated interrupt while finalising the archive") + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + # 只在首次导出注入故障，之后要验证同一实例可以继续工作 + with monkeypatch.context() as failing_finalisation: + failing_finalisation.setattr( + zipfile.ZipFile, "_write_end_record", interrupt_end_record + ) + with pytest.raises(_TeardownInterrupt): + await exporter.export_all(output_dir=str(backup_dir)) + + # 执行器必须已释放：这正是再入保护检查的状态 + assert exporter._archive_executor is None + + # 同一实例仍然可用，而不是从此永远抛 "already running" + zip_path = await exporter.export_all(output_dir=str(backup_dir)) + with zipfile.ZipFile(zip_path) as zf: + assert zf.testzip() is None + + @pytest.mark.asyncio + async def test_successful_export_still_returns_usable_archive( + self, tmp_path, monkeypatch + ): + """成功路径不受影响：仍然返回路径，且归档完整可读""" + data_dir = self._make_data_dir(tmp_path) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(data_dir / "plugins")}, + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + zip_path = await exporter.export_all(output_dir=str(backup_dir)) + + assert os.path.exists(zip_path) + with zipfile.ZipFile(zip_path) as zf: + assert zf.testzip() is None + assert "manifest.json" in zf.namelist() + + # 导出器释放了归档执行器：同一实例可以立刻再导出一份完整备份 + second_zip_path = await exporter.export_all(output_dir=str(backup_dir)) + assert os.path.exists(second_zip_path) + with zipfile.ZipFile(second_zip_path) as zf: + assert zf.testzip() is None + assert "manifest.json" in zf.namelist() + + @pytest.mark.asyncio + async def test_real_error_wins_over_cancellation_during_teardown( + self, tmp_path, monkeypatch + ): + """收尾期间被取消时，真正的导出错误必须胜出 + + 修复前收尾无条件重抛 CancelledError：调用方的 ``except Exception`` 看不到 + 真正的失败原因，任务会一直停在 processing。 + """ + data_dir = self._make_data_dir(tmp_path) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(data_dir / "plugins")}, + ) + + # 本用例需要收尾的 close 停在归档线程里，以便在收尾期间投递取消 + close_started = threading.Event() + release_close = threading.Event() + real_close = zipfile.ZipFile.close + + # 记录本次导出自己的归档实例。按位置取出（当且仅当本次导出真的调用过它， + # 出口才有值），因此取值前后都不需要额外的时序假设。 + archive_holder: list[zipfile.ZipFile] = [] + + def boom(zf): + # 这里拿到的就是本次导出自己打开的归档实例，先记下来 + # （即收尾 close 的同步时机），再抛出真正的导出错误 + archive_holder.append(zf) + raise RuntimeError("simulated directory export failure") + + # 同步必须锚定「本次导出的那个归档实例」：``close`` 是类级补丁，而 + # ``ZipFile.__del__`` 同样会调用它——同批次先跑完的用例留下的归档一旦被 gc + # 回收，就会替真正的收尾把 ``close_started`` 提前置位（实测这正是修复前 + # ~67s = 两个 30s 超时的来源：一次陈旧 close、一次真正的 close）。取消因此 + # 打在 ``await asyncio.shield(open_task)`` 这个「打开」等待上，而那时还没有 + # 任何导出错误，收尾只能重抛 CancelledError，用例随之失败。 + # 因此这里只对本次导出的归档阻塞并置位，其余实例一律直接透传给真实 close， + # 不发信号也不阻塞：同步点唯一，且不依赖 gc 时机或用例顺序。 + def blocking_close(zf_self): + if not archive_holder or zf_self is not archive_holder[0]: + return real_close(zf_self) + close_started.set() + release_close.wait(timeout=10) + return real_close(zf_self) + + monkeypatch.setattr(zipfile.ZipFile, "close", blocking_close) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + # 只在首次导出注入故障，之后要验证同一实例可以继续工作 + with monkeypatch.context() as failing_export: + failing_export.setattr( + "astrbot.core.backup.exporter._write_backup_directories", boom + ) + + export_task = asyncio.create_task( + exporter.export_all(output_dir=str(backup_dir)) + ) + try: + # 本地临时目录导出只需数秒，10s 足够；不再各花 30s 死等 + assert await asyncio.to_thread(close_started.wait, 10) + export_task.cancel() + finally: + release_close.set() + + # 真正的失败必须原样上报，而不是被收尾吸收的取消顶替 + with pytest.raises(RuntimeError, match="simulated directory export failure"): + await export_task + + # 残留 ZIP 必须被清理 + assert list(backup_dir.glob("*.zip")) == [] + + # 执行器必须已释放：这正是再入保护检查的状态 + assert exporter._archive_executor is None + + # 可观测证据：同一实例的后续导出立刻成功，而不是永远抛 "already running" + zip_path = await exporter.export_all(output_dir=str(backup_dir)) + with zipfile.ZipFile(zip_path) as zf: + assert zf.testzip() is None + assert "manifest.json" in zf.namelist()