From 377635738d57b8b7ac4b64223e0a094e52ddcf14 Mon Sep 17 00:00:00 2001 From: JosephTian876 Date: Wed, 7 Oct 2026 01:27:51 +0800 Subject: [PATCH 1/4] fix: prevent backup export from blocking the event loop export_all() ran every ZIP compression and directory walk on the asyncio event loop thread. On a 2.4 GB dataset that froze the whole bot for 76 s, and AstrBot's own diagnostic recorded a 119.578 s loop lag, well past its 15 s threshold. While the loop is frozen the reverse WebSocket to the OneBot implementation cannot answer, so the implementation times out after ~90 s and reconnects. Archive work now runs on a dedicated single-worker ThreadPoolExecutor owned by export_all. A single worker is used rather than the shared default pool because zipfile.ZipFile is not safe for concurrent use, and it doubles as a FIFO queue: the archive is opened and closed on that worker too, so a close can never overlap an in-flight write. A failed finalisation is reported as a failure instead of returning a path to an unusable archive. Backup format, ZIP layout, manifest.json, checksums, progress_callback semantics and every public signature are unchanged. --- astrbot/core/backup/exporter.py | 577 +++++++++++++++++++++++--------- 1 file changed, 413 insertions(+), 164 deletions(-) diff --git a/astrbot/core/backup/exporter.py b/astrbot/core/backup/exporter.py index 2bd0c30178b..9afa1b71a62 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,122 +145,278 @@ 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 + ) + await self._run_in_archive_thread( + zf.writestr, "databases/kb_metadata.json", kb_meta_json ) - zf.writestr("databases/main_db.json", main_db_json) - self._add_checksum("databases/main_db.json", main_db_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) - mcp_config_path = Path(get_astrbot_data_path()) / "mcp_server.json" - if mcp_config_path.is_file(): - mcp_content = mcp_config_path.read_bytes() - zf.writestr("config/mcp_server.json", mcp_content) - self._add_checksum("config/mcp_server.json", mcp_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) + mcp_config_path = Path(get_astrbot_data_path()) / "mcp_server.json" + if mcp_config_path.is_file(): + mcp_content = mcp_config_path.read_bytes() + await self._run_in_archive_thread( + zf.writestr, "config/mcp_server.json", mcp_content + ) + self._add_checksum("config/mcp_server.json", mcp_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]]: """导出主数据库所有表""" @@ -285,7 +495,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}") @@ -294,87 +506,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 实例转换为字典 @@ -480,3 +631,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}") From e130525e49a8f076429fc37f35d5d5484bc075f9 Mon Sep 17 00:00:00 2001 From: JosephTian876 Date: Wed, 7 Oct 2026 01:27:51 +0800 Subject: [PATCH 2/4] test: add regression tests for non-blocking backup export Cover the event-loop offload, archive structural stability, cancellation teardown, re-entrancy, and the rule that a failed finalisation must not be reported as a successful export. Each guard fails on the previous implementation. --- tests/test_backup.py | 664 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 664 insertions(+) diff --git a/tests/test_backup.py b/tests/test_backup.py index 7fbd93a50ed..24522e4a3c4 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -1,9 +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 @@ -1741,3 +1745,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() From c6bd743c8c76b8a7c9923e006bfc7ede0fa3f173 Mon Sep 17 00:00:00 2001 From: JosephTian876 <204833675+JosephTian876@users.noreply.github.com> Date: Sun, 18 Oct 2026 03:39:00 +0800 Subject: [PATCH 3/4] fix: move serialization off the loop and clean up failed archives Address the review findings on the non-blocking backup export: - JSON serialization (main DB, KB metadata, KB documents, manifest) and manifest generation now run on the archive worker instead of the event loop. The manifest walks every KB media directory, so it could freeze the loop just as the compression used to. - Clear self._checksums at the start of every export, so a reused exporter cannot carry the previous run's checksums into a new manifest that no longer contains those files. - Capture BaseException from the archive open and close paths into open_error/close_error instead of letting it escape the teardown and skip the partial-ZIP removal. A failed open with no other error is now reported instead of returning a path to a file that was never written. Each guard has a regression test that fails on the previous revision. --- astrbot/core/backup/exporter.py | 75 ++++++++++++----- tests/test_backup.py | 140 +++++++++++++++++++++++++++++++- 2 files changed, 194 insertions(+), 21 deletions(-) diff --git a/astrbot/core/backup/exporter.py b/astrbot/core/backup/exporter.py index 9afa1b71a62..89ba73dc21c 100644 --- a/astrbot/core/backup/exporter.py +++ b/astrbot/core/backup/exporter.py @@ -73,10 +73,16 @@ def __init__( 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. + """Run a blocking export operation on the executor that owns the archive. + + The executor also carries JSON serialization and manifest generation, so + neither of them can run on the event loop while a large export is being + produced. Because it has exactly one worker, those CPU-bound calls and the + archive writes are serialised against each other, which is what makes the + shared ``ZipFile`` safe to use. Args: - func: The blocking archive callable to execute, e.g. ``zf.write``. + func: The blocking callable to execute, e.g. ``zf.write`` or ``json.dumps``. *args: Positional arguments forwarded to ``func``. **kwargs: Keyword arguments forwarded to ``func``. @@ -84,8 +90,8 @@ async def _run_in_archive_thread( Whatever ``func`` returns. Raises: - RuntimeError: If the private archive helpers are called outside - ``export_all``, i.e. while no archive executor owns the archive. + RuntimeError: If the private export helpers are called outside + ``export_all``, i.e. while no archive executor owns the export. """ # 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), @@ -133,6 +139,11 @@ async def export_all( "separate AstrBotExporter for each concurrent export." ) + # Every export builds its own manifest. Without the reset, a reused exporter + # carries the previous run's checksums into the new manifest, which then lists + # files that are absent from the new ZIP. + self._checksums.clear() + if output_dir is None: output_dir = get_astrbot_backups_path() @@ -160,6 +171,7 @@ async def export_all( # the teardown can tell "the archive was never finished" apart from "the archive # finished but could not be finalised". export_error: BaseException | None = None + open_error: BaseException | None = None close_error: BaseException | None = None cancelled_while_cleaning = False @@ -183,8 +195,10 @@ async def export_all( 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 + # Serialization of a large database is CPU-bound work that used to run on the + # event loop; the archive worker owns it now, together with the ZIP writes. + main_db_json = await self._run_in_archive_thread( + 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 @@ -205,8 +219,8 @@ async def export_all( "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 + kb_meta_json = await self._run_in_archive_thread( + json.dumps, kb_meta_data, ensure_ascii=False, indent=2, default=str ) await self._run_in_archive_thread( zf.writestr, "databases/kb_metadata.json", kb_meta_json @@ -229,8 +243,8 @@ async def export_all( 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_json = await self._run_in_archive_thread( + 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) @@ -286,8 +300,14 @@ async def export_all( # 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) + # The manifest walks every knowledge-base media directory, so it is generated + # and serialized on the archive worker as well instead of on the event loop. + manifest = await self._run_in_archive_thread( + self._generate_manifest, main_data, kb_meta_data, dir_stats + ) + manifest_json = await self._run_in_archive_thread( + json.dumps, manifest, ensure_ascii=False, indent=2 + ) await self._run_in_archive_thread( zf.writestr, "manifest.json", manifest_json ) @@ -325,9 +345,11 @@ async def export_all( 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. + except BaseException: + # The open itself failed, including a BaseException such as + # KeyboardInterrupt raised on the archive worker. Its result is + # read below to report why, so swallow it here and let the + # teardown continue instead of skipping the ZIP cleanup. pass if not open_task.cancelled(): try: @@ -336,7 +358,10 @@ async def export_all( # the instance so its handle gets closed instead of leaked. if zf is None: zf = open_task.result() - except Exception as e: + except BaseException as e: + # A BaseException escaping here would otherwise jump over the + # partial-file removal below, leaving an unusable ZIP behind. + open_error = e logger.warning(f"打开备份归档失败: {e}") if zf is not None: @@ -354,8 +379,10 @@ async def export_all( await asyncio.shield(close_task) except asyncio.CancelledError: cancelled_while_cleaning = True - except Exception: - # A failed close is an export failure, reported below. + except BaseException: + # A failed close is an export failure, reported below. This + # also covers a BaseException such as KeyboardInterrupt, which + # would otherwise skip the partial-file removal. pass if not close_task.cancelled(): try: @@ -374,12 +401,13 @@ async def export_all( if ( export_error is not None + or open_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 + # was cancelled, or the archive could not be opened or finalised and is + # left in an unusable state. 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 @@ -391,6 +419,13 @@ async def export_all( except OSError as remove_error: logger.warning(f"清理残缺备份文件失败: {remove_error}") + if open_error is not None and export_error is None: + # The archive could not be created and no other error is propagating. This + # is defensive: the failed open normally raises at the await above, so this + # branch only exists to guarantee export_all never returns a path to a file + # that was never produced. + raise open_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 diff --git a/tests/test_backup.py b/tests/test_backup.py index 24522e4a3c4..9790860253e 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -11,6 +11,7 @@ import zipfile from datetime import datetime from pathlib import Path +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -260,6 +261,50 @@ async def test_export_all_creates_zip( assert "databases/main_db.json" in namelist assert "config/cmd_config.json" in namelist + @pytest.mark.asyncio + async def test_export_manifest_drops_stale_checksums(self, tmp_path, monkeypatch): + """复用同一实例导出两次时,manifest 不得残留上一轮的校验和""" + 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)}, + ) + + 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), + ) + exporter = AstrBotExporter( + main_db=mock_db, + kb_manager=None, + config_path=str(data_dir / "cmd_config.json"), + ) + backup_dir = tmp_path / "backups" + + await exporter.export_all(output_dir=str(backup_dir)) + assert "config/cmd_config.json" in exporter._checksums + + # 第二轮不再导出配置文件,manifest 也就不能继续列出它的校验和 + os.remove(data_dir / "cmd_config.json") + second_zip = await exporter.export_all(output_dir=str(backup_dir)) + + with zipfile.ZipFile(second_zip) as zf: + assert "config/cmd_config.json" not in zf.namelist() + manifest = json.loads(zf.read("manifest.json")) + assert "config/cmd_config.json" not in manifest["checksums"] + for path in manifest["checksums"]: + assert path in zf.namelist() + class TestAstrBotImporter: """AstrBotImporter 类测试""" @@ -1909,6 +1954,64 @@ async def test_event_loop_stays_responsive_during_export( max_gap = max(gaps) assert max_gap < 0.5, f"事件循环最长停顿 {max_gap:.3f}s" + @pytest.mark.asyncio + async def test_serialization_and_manifest_leave_event_loop_thread( + self, tmp_path, monkeypatch + ): + """JSON 序列化与 manifest 生成也必须离开事件循环线程 + + 修复前这两步在事件循环上同步执行:数据库很大或知识库媒体目录很大时, + 即使 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)}, + ) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + serialization_threads: set[int] = set() + manifest_threads: set[int] = set() + real_dumps = json.dumps + real_generate_manifest = AstrBotExporter._generate_manifest + + def spy_dumps(*args, **kwargs): + serialization_threads.add(threading.get_ident()) + return real_dumps(*args, **kwargs) + + def spy_generate_manifest(exporter_self, *args, **kwargs): + manifest_threads.add(threading.get_ident()) + return real_generate_manifest(exporter_self, *args, **kwargs) + + # 只替换 exporter 模块对 json 的引用,避免影响同进程其他 json 调用 + from astrbot.core.backup import exporter as exporter_module + + monkeypatch.setattr(exporter_module, "json", SimpleNamespace(dumps=spy_dumps)) + monkeypatch.setattr( + AstrBotExporter, "_generate_manifest", spy_generate_manifest + ) + + loop_thread_id = threading.get_ident() + zip_path = await exporter.export_all(output_dir=str(backup_dir)) + + assert os.path.exists(zip_path) + # 两步都确实执行过,且都没有落在事件循环线程上 + assert serialization_threads + assert loop_thread_id not in serialization_threads + assert manifest_threads + assert loop_thread_id not in manifest_threads + + with zipfile.ZipFile(zip_path, "r") as zf: + assert "manifest.json" in zf.namelist() + @pytest.mark.asyncio async def test_export_failure_shuts_down_executor_and_cleans_zip( self, tmp_path, monkeypatch @@ -2290,12 +2393,45 @@ def interrupt_end_record(*args, **kwargs): # 执行器必须已释放:这正是再入保护检查的状态 assert exporter._archive_executor is None + # BaseException 不得越过残缺 ZIP 的清理:收尾失败后磁盘上不能留下坏归档 + assert list(backup_dir.glob("*.zip")) == [] # 同一实例仍然可用,而不是从此永远抛 "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_base_exception_during_open_still_removes_partial_zip( + self, tmp_path, monkeypatch + ): + """打开归档时抛 BaseException 也必须清理残留 ZIP 并释放执行器 + + 修复前收尾只捕获 Exception:BaseException 会直接从 ``open_task.result()`` + 冲出 finally,跳过残缺文件清理,把一个永远无法打开的 ZIP 留在备份目录里。 + """ + data_dir = self._make_data_dir(tmp_path) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {"plugins": str(data_dir / "plugins")}, + ) + + def exploding_zipfile(*args, **kwargs): + # 先落盘一个残缺文件,模拟真实失败后磁盘上的残留 + Path(args[0]).write_bytes(b"PK\x03\x04partial-archive") + raise _TeardownInterrupt("simulated interrupt while opening the archive") + + monkeypatch.setattr(zipfile, "ZipFile", exploding_zipfile) + + exporter = self._make_exporter(data_dir) + backup_dir = tmp_path / "backups" + + with pytest.raises(_TeardownInterrupt): + await exporter.export_all(output_dir=str(backup_dir)) + + assert list(backup_dir.glob("*.zip")) == [] + assert exporter._archive_executor is None + @pytest.mark.asyncio async def test_successful_export_still_returns_usable_archive( self, tmp_path, monkeypatch @@ -2391,7 +2527,9 @@ def blocking_close(zf_self): release_close.set() # 真正的失败必须原样上报,而不是被收尾吸收的取消顶替 - with pytest.raises(RuntimeError, match="simulated directory export failure"): + with pytest.raises( + RuntimeError, match="simulated directory export failure" + ): await export_task # 残留 ZIP 必须被清理 From 49587979e731a803da8544513a729bf126b62605 Mon Sep 17 00:00:00 2001 From: JosephTian876 Date: Mon, 26 Oct 2026 05:47:10 +0800 Subject: [PATCH 4/4] test: decouple archive structure test from the shared data directory The exporter now also archives `mcp_server.json` from the AstrBot data directory. The archive structure test asserts an exact member list while the exporter still read that file from the process data root, so the assertion only held on a checkout without one. Point `get_astrbot_data_path` at the temporary data directory used by the test, matching how the neighbouring `get_backup_directories` patch already isolates the archive from real user data. --- tests/test_backup.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_backup.py b/tests/test_backup.py index 9790860253e..dab08d6a82a 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -2194,6 +2194,12 @@ async def test_export_structure_is_stable(self, tmp_path, monkeypatch): "astrbot.core.backup.exporter.get_backup_directories", lambda: {"plugins": str(plugin_dir)}, ) + # 导出器还会归档数据目录下的 mcp_server.json;把它指向临时目录, + # 成员清单断言才不会取决于检出目录里是否存在该文件。 + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_astrbot_data_path", + lambda: str(data_dir), + ) first_dir = tmp_path / "backups-a" second_dir = tmp_path / "backups-b"