refactor: 抽离应用 runner 并优化退出清理

NanmiCoder · NanmiCoder · commit 8a0fd49b960e · 2025-12-15T18:06:57.000+08:00
- 新增 tools/app_runner.py 统一信号/取消/清理超时逻辑
- main.py 精简为业务入口与资源清理实现
- CDPBrowserManager 不再覆盖已有 SIGINT/SIGTERM 处理器
diff --git a/main.py b/main.py
@@ -19,9 +19,7 @@
 
 
 import asyncio
-import sys
-import signal
-from typing import Optional
+from typing import Optional, Type
 
 import cmd_arg
 import config
@@ -39,7 +37,7 @@
 
 
 class CrawlerFactory:
-    CRAWLERS = {
+    CRAWLERS: dict[str, Type[AbstractCrawler]] = {
         "xhs": XiaoHongShuCrawler,
         "dy": DouYinCrawler,
         "ks": KuaishouCrawler,
@@ -53,115 +51,96 @@ class CrawlerFactory:
     def create_crawler(platform: str) -> AbstractCrawler:
         crawler_class = CrawlerFactory.CRAWLERS.get(platform)
         if not crawler_class:
-            raise ValueError(
-                "Invalid Media Platform Currently only supported xhs or dy or ks or bili ..."
-            )
+            supported = ", ".join(sorted(CrawlerFactory.CRAWLERS))
+            raise ValueError(f"Invalid media platform: {platform!r}. Supported: {supported}")
         return crawler_class()
 
 
 crawler: Optional[AbstractCrawler] = None
 
 
-# persist-1<persist1@126.com>
-# 原因：增加 --init_db 功能，用于数据库初始化。
-# 副作用：无
-# 回滚策略：还原此文件。
-async def main():
-    # Init crawler
+def _flush_excel_if_needed() -> None:
+    if config.SAVE_DATA_OPTION != "excel":
+        return
+
+    try:
+        from store.excel_store_base import ExcelStoreBase
+
+        ExcelStoreBase.flush_all()
+        print("[Main] Excel files saved successfully")
+    except Exception as e:
+        print(f"[Main] Error flushing Excel data: {e}")
+
+
+async def _generate_wordcloud_if_needed() -> None:
+    if config.SAVE_DATA_OPTION != "json" or not config.ENABLE_GET_WORDCLOUD:
+        return
+
+    try:
+        file_writer = AsyncFileWriter(
+            platform=config.PLATFORM,
+            crawler_type=crawler_type_var.get(),
+        )
+        await file_writer.generate_wordcloud_from_comments()
+    except Exception as e:
+        print(f"[Main] Error generating wordcloud: {e}")
+
+
+async def main() -> None:
     global crawler
 
-    # parse cmd
     args = await cmd_arg.parse_cmd()
-
-    # init db
     if args.init_db:
         await db.init_db(args.init_db)
         print(f"Database {args.init_db} initialized successfully.")
-        return  # Exit the main function cleanly
-
-
+        return
 
     crawler = CrawlerFactory.create_crawler(platform=config.PLATFORM)
     await crawler.start()
 
-    # Flush Excel data if using Excel export
-    if config.SAVE_DATA_OPTION == "excel":
-        try:
-            from store.excel_store_base import ExcelStoreBase
-            ExcelStoreBase.flush_all()
-            print("[Main] Excel files saved successfully")
-        except Exception as e:
-            print(f"[Main] Error flushing Excel data: {e}")
+    _flush_excel_if_needed()
 
     # Generate wordcloud after crawling is complete
     # Only for JSON save mode
-    if config.SAVE_DATA_OPTION == "json" and config.ENABLE_GET_WORDCLOUD:
-        try:
-            file_writer = AsyncFileWriter(
-                platform=config.PLATFORM,
-                crawler_type=crawler_type_var.get()
-            )
-            await file_writer.generate_wordcloud_from_comments()
-        except Exception as e:
-            print(f"Error generating wordcloud: {e}")
+    await _generate_wordcloud_if_needed()
 
 
-async def async_cleanup():
-    """异步清理函数，用于处理CDP浏览器等异步资源"""
+async def async_cleanup() -> None:
     global crawler
     if crawler:
-        # 检查并清理CDP浏览器
-        if hasattr(crawler, 'cdp_manager') and crawler.cdp_manager:
+        if getattr(crawler, "cdp_manager", None):
             try:
-                await crawler.cdp_manager.cleanup(force=True)  # 强制清理浏览器进程
+                await crawler.cdp_manager.cleanup(force=True)
             except Exception as e:
-                # 只在非预期错误时打印
                 error_msg = str(e).lower()
                 if "closed" not in error_msg and "disconnected" not in error_msg:
                     print(f"[Main] 清理CDP浏览器时出错: {e}")
 
-        # 检查并清理标准浏览器上下文（仅在非CDP模式下）
-        elif hasattr(crawler, 'browser_context') and crawler.browser_context:
+        elif getattr(crawler, "browser_context", None):
             try:
-                # 检查上下文是否仍然打开
-                if hasattr(crawler.browser_context, 'pages'):
-                    await crawler.browser_context.close()
+                await crawler.browser_context.close()
             except Exception as e:
-                # 只在非预期错误时打印
                 error_msg = str(e).lower()
                 if "closed" not in error_msg and "disconnected" not in error_msg:
                     print(f"[Main] 关闭浏览器上下文时出错: {e}")
 
-    # 关闭数据库连接
-    if config.SAVE_DATA_OPTION in ["db", "sqlite"]:
+    if config.SAVE_DATA_OPTION in ("db", "sqlite"):
         await db.close()
 
-def cleanup():
-    """同步清理函数"""
-    try:
-        # 创建新的事件循环来执行异步清理
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        loop.run_until_complete(async_cleanup())
-        loop.close()
-    except Exception as e:
-        print(f"[Main] 清理时出错: {e}")
-
-
-def signal_handler(signum, _frame):
-    """信号处理器，处理Ctrl+C等中断信号"""
-    print(f"\n[Main] 收到中断信号 {signum}，正在清理资源...")
-    cleanup()
-    sys.exit(0)
-
 if __name__ == "__main__":
-    # 注册信号处理器
-    signal.signal(signal.SIGINT, signal_handler)   # Ctrl+C
-    signal.signal(signal.SIGTERM, signal_handler)  # 终止信号
+    from tools.app_runner import run
+
+    def _force_stop() -> None:
+        c = crawler
+        if not c:
+            return
+        cdp_manager = getattr(c, "cdp_manager", None)
+        launcher = getattr(cdp_manager, "launcher", None)
+        if not launcher:
+            return
+        try:
+            launcher.cleanup()
+        except Exception:
+            pass
 
-    try:
-        asyncio.get_event_loop().run_until_complete(main())
-    except KeyboardInterrupt:
-        print("\n[Main] 收到键盘中断，正在清理资源...")
-    finally:
-        cleanup()
+    run(main, async_cleanup, cleanup_timeout_seconds=15.0, on_first_interrupt=_force_stop)
diff --git a/tools/app_runner.py b/tools/app_runner.py
@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2025 relakkes@gmail.com
+#
+# This file is part of MediaCrawler project.
+# Repository: https://github.com/NanmiCoder/MediaCrawler
+# GitHub: https://github.com/NanmiCoder
+# Licensed under NON-COMMERCIAL LEARNING LICENSE 1.1
+#
+
+# 声明：本代码仅供学习和研究目的使用。使用者应遵守以下原则：
+# 1. 不得用于任何商业用途。
+# 2. 使用时应遵守目标平台的使用条款和robots.txt规则。
+# 3. 不得进行大规模爬取或对平台造成运营干扰。
+# 4. 应合理控制请求频率，避免给目标平台带来不必要的负担。
+# 5. 不得用于任何非法或不当的用途。
+#
+# 详细许可条款请参阅项目根目录下的LICENSE文件。
+# 使用本代码即表示您同意遵守上述原则和LICENSE中的所有条款。
+
+
+from __future__ import annotations
+
+import asyncio
+import os
+import signal
+from collections.abc import Awaitable, Callable
+from typing import Optional
+
+AsyncFn = Callable[[], Awaitable[None]]
+
+
+def run(
+    app_main: AsyncFn,
+    app_cleanup: AsyncFn,
+    *,
+    cleanup_timeout_seconds: float = 15.0,
+    on_first_interrupt: Optional[Callable[[], None]] = None,
+    force_exit_code: int = 130,
+) -> None:
+    async def _cleanup_with_timeout() -> None:
+        try:
+            await asyncio.wait_for(asyncio.shield(app_cleanup()), timeout=cleanup_timeout_seconds)
+        except asyncio.TimeoutError:
+            print(f"[Main] 清理超时({cleanup_timeout_seconds}s)，跳过剩余清理。")
+
+    async def _cancel_remaining_tasks(timeout_seconds: float = 2.0) -> None:
+        current = asyncio.current_task()
+        tasks = [t for t in asyncio.all_tasks() if t is not current and not t.done()]
+        if not tasks:
+            return
+        for t in tasks:
+            t.cancel()
+        try:
+            await asyncio.wait_for(
+                asyncio.gather(*tasks, return_exceptions=True),
+                timeout=timeout_seconds,
+            )
+        except asyncio.TimeoutError:
+            pass
+
+    async def _runner() -> None:
+        loop = asyncio.get_running_loop()
+        runner_task = asyncio.current_task()
+        if runner_task is None:
+            raise RuntimeError("Runner task not found")
+
+        shutdown_requested = False
+
+        def _on_signal(signum: int) -> None:
+            nonlocal shutdown_requested
+
+            if shutdown_requested:
+                print("[Main] 再次收到中断信号，强制退出。")
+                os._exit(force_exit_code)
+
+            shutdown_requested = True
+            print(f"\n[Main] 收到中断信号 {signum}，正在退出(清理最多{cleanup_timeout_seconds}s)...")
+
+            if on_first_interrupt is not None:
+                try:
+                    on_first_interrupt()
+                except Exception:
+                    pass
+
+            runner_task.cancel()
+
+        try:
+            loop.add_signal_handler(signal.SIGINT, _on_signal, signal.SIGINT)
+            loop.add_signal_handler(signal.SIGTERM, _on_signal, signal.SIGTERM)
+        except NotImplementedError:
+            signal.signal(signal.SIGINT, lambda signum, _frame: _on_signal(signum))
+            signal.signal(signal.SIGTERM, lambda signum, _frame: _on_signal(signum))
+
+        cancelled = False
+        try:
+            await app_main()
+        except asyncio.CancelledError:
+            cancelled = True
+        finally:
+            try:
+                await _cleanup_with_timeout()
+            except Exception as e:
+                print(f"[Main] 清理时出错: {e}")
+            await _cancel_remaining_tasks()
+
+        if cancelled:
+            return
+
+    asyncio.run(_runner())
diff --git a/tools/cdp_browser.py b/tools/cdp_browser.py
@@ -60,19 +60,36 @@ def sync_cleanup():
         # 注册atexit清理
         atexit.register(sync_cleanup)
 
-        # 注册信号处理器
+        # 注册信号处理器（仅在没有自定义处理器时注册，避免覆盖主入口的信号处理逻辑）
+        prev_sigint = signal.getsignal(signal.SIGINT)
+        prev_sigterm = signal.getsignal(signal.SIGTERM)
+
         def signal_handler(signum, frame):
             """信号处理器"""
             utils.logger.info(f"[CDPBrowserManager] 收到信号 {signum}，清理浏览器进程")
             if self.launcher and self.launcher.browser_process:
                 self.launcher.cleanup()
-            # 重新引发KeyboardInterrupt以便正常退出流程
+
             if signum == signal.SIGINT:
+                if prev_sigint == signal.default_int_handler:
+                    return prev_sigint(signum, frame)
                 raise KeyboardInterrupt
 
+            raise SystemExit(0)
+
+        install_sigint = prev_sigint in (signal.default_int_handler, signal.SIG_DFL)
+        install_sigterm = prev_sigterm == signal.SIG_DFL
+
         # 注册SIGINT (Ctrl+C) 和 SIGTERM
-        signal.signal(signal.SIGINT, signal_handler)
-        signal.signal(signal.SIGTERM, signal_handler)
+        if install_sigint:
+            signal.signal(signal.SIGINT, signal_handler)
+        else:
+            utils.logger.info("[CDPBrowserManager] 已存在SIGINT处理器，跳过注册以避免覆盖")
+
+        if install_sigterm:
+            signal.signal(signal.SIGTERM, signal_handler)
+        else:
+            utils.logger.info("[CDPBrowserManager] 已存在SIGTERM处理器，跳过注册以避免覆盖")
 
         self._cleanup_registered = True
         utils.logger.info("[CDPBrowserManager] 清理处理器已注册")
