diff --git a/api_browser.py b/api_browser.py index e3ff2f9..5dfffcc 100755 --- a/api_browser.py +++ b/api_browser.py @@ -9,6 +9,7 @@ import requests from bs4 import BeautifulSoup import re import time +import atexit from typing import Optional, Callable from dataclasses import dataclass @@ -38,6 +39,8 @@ class APIBrowser: self.logged_in = False self.log_callback = log_callback self.stop_flag = False + self._closed = False # 防止重复关闭 + # 设置代理 if proxy_config and proxy_config.get("server"): proxy_server = proxy_config["server"] @@ -49,6 +52,9 @@ class APIBrowser: else: self.proxy_server = None + # 注册退出清理函数 + atexit.register(self._cleanup_on_exit) + def log(self, message: str): """记录日志""" if self.log_callback: @@ -382,7 +388,29 @@ class APIBrowser: def close(self): """关闭会话""" + if self._closed: + return + self._closed = True + try: self.session.close() except: pass + + def _cleanup_on_exit(self): + """进程退出时的清理函数（由atexit调用）""" + if not self._closed: + try: + self.session.close() + self._closed = True + except: + pass + + def __enter__(self): + """Context manager支持 - 进入""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager支持 - 退出""" + self.close() + return False # 不抑制异常 diff --git a/app.py b/app.py index f845395..77c771d 100755 --- a/app.py +++ b/app.py @@ -3484,7 +3484,67 @@ def batch_stop_accounts(): "stopped": stopped }) +def cleanup_on_exit(): + """应用退出时的清理函数""" + import signal + import sys + + print("\n正在清理资源...") + + # 1. 停止所有运行中的任务 + print("- 停止运行中的任务...") + for account_id, thread in list(active_tasks.items()): + try: + # 设置停止标志 + if account_id in user_accounts: + for user_id in user_accounts: + if account_id in user_accounts[user_id]: + user_accounts[user_id][account_id].should_stop = True + except: + pass + + # 2. 等待所有线程完成（最多等待5秒） + print("- 等待线程退出...") + for account_id, thread in list(active_tasks.items()): + try: + if thread and thread.is_alive(): + thread.join(timeout=2) + except: + pass + + # 3. 关闭浏览器工作线程池 + print("- 关闭浏览器线程池...") + try: + shutdown_browser_worker_pool() + except: + pass + + # 4. 关闭数据库连接池 + print("- 关闭数据库连接池...") + try: + db_pool._pool.close_all() if db_pool._pool else None + except: + pass + + print("✓ 资源清理完成") + + if __name__ == '__main__': + import signal + import atexit + + # 注册退出清理函数 + atexit.register(cleanup_on_exit) + + # 注册信号处理器 + def signal_handler(sig, frame): + print("\n\n收到退出信号，正在关闭...") + cleanup_on_exit() + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + print("=" * 60) print("知识管理平台自动化工具 - 多用户版") print("=" * 60) diff --git a/playwright_automation.py b/playwright_automation.py index 712680b..8b9af25 100755 --- a/playwright_automation.py +++ b/playwright_automation.py @@ -11,6 +11,7 @@ from playwright.sync_api import sync_playwright, Browser, BrowserContext, Page, import time import json import threading +import atexit from typing import Optional, Callable from dataclasses import dataclass from app_config import get_config @@ -147,6 +148,10 @@ class PlaywrightAutomation: self.context: Optional[BrowserContext] = None self.page: Optional[Page] = None self.main_page: Optional[Page] = None + self._closed = False # 防止重复关闭 + + # 注册退出清理函数，确保进程异常退出时也能关闭浏览器 + atexit.register(self._cleanup_on_exit) def log(self, message: str): """记录日志""" @@ -1269,7 +1274,16 @@ class PlaywrightAutomation: Returns: 是否截图成功 """ + import os + temp_filepath = None + try: + # 确保目录存在 + os.makedirs(os.path.dirname(filepath), exist_ok=True) + + # 先保存到临时文件 + temp_filepath = filepath + '.tmp' + # 使用最高质量设置截图 # type='jpeg' 指定JPEG格式（支持quality参数） # quality=100 表示100%的JPEG质量（范围0-100，最高质量） @@ -1277,19 +1291,47 @@ class PlaywrightAutomation: # 视口分辨率 2560x1440 确保高清晰度 # 这样可以生成更清晰的截图，大小约500KB-1MB左右 self.main_page.screenshot( - path=filepath, + path=temp_filepath, type='jpeg', full_page=True, quality=100 ) - self.log(f"截图已保存: {filepath}") + + # 验证文件是否成功创建且大小合理 + if not os.path.exists(temp_filepath): + raise FileNotFoundError(f"截图文件未创建: {temp_filepath}") + + file_size = os.path.getsize(temp_filepath) + if file_size == 0: + raise ValueError(f"截图文件为空: {temp_filepath}") + + # 文件验证通过，重命名为最终文件名 + if os.path.exists(filepath): + os.remove(filepath) # 删除旧文件 + os.rename(temp_filepath, filepath) + + self.log(f"截图已保存: {filepath} ({file_size} bytes)") return True + except Exception as e: self.log(f"截图失败: {str(e)}") + + # 清理临时文件 + try: + if temp_filepath and os.path.exists(temp_filepath): + os.remove(temp_filepath) + except Exception as cleanup_error: + self.log(f"清理临时文件失败: {cleanup_error}") + return False def close(self): """完全关闭浏览器进程(每个账号独立)并确保资源释放""" + # 防止重复关闭 + if self._closed: + return + self._closed = True + errors = [] # 第一步:关闭上下文 @@ -1339,6 +1381,30 @@ class PlaywrightAutomation: self.log(f"资源清理完成,但有{len(errors)}个警告") # else部分日志已精简 + def _cleanup_on_exit(self): + """进程退出时的清理函数（由atexit调用）""" + if not self._closed: + try: + # 静默关闭，避免在退出时产生过多日志 + if self.context: + self.context.close() + if self.browser: + self.browser.close() + if self.playwright: + self.playwright.stop() + self._closed = True + except: + pass # 退出时忽略所有错误 + + def __enter__(self): + """Context manager支持 - 进入""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager支持 - 退出""" + self.close() + return False # 不抑制异常 + # 简单的测试函数 if __name__ == "__main__":