import csv import json import os import threading import time import ctypes from concurrent.futures import ThreadPoolExecutor from datetime import datetime from config import ( channel_name, DASHBOARD_HOST, DASHBOARD_PORT, CHANNEL_NAMESPACE, FINAL_REPORT_CSV, INSTANCE_NAME, REDIS_HOST, REDIS_DB, REDIS_PORT, REPORT_DIR, WORKER_REPORT_CSV, WORKER_STALE_TIMEOUT, WORKER_STATUS_PUSH_INTERVAL, APK_DOWNLOAD_QUEUE_INTERVAL, APK_PREFETCH_POLL_INTERVAL, ) from dashboard_main import run_dashboard from log_manager import logger from redis_task_distribute import RedisTaskDispatcher def _set_console_title(): title = ( f"AUTOOL Dispatcher [{INSTANCE_NAME}] " f"Redis {REDIS_PORT}/{REDIS_DB} Dashboard {DASHBOARD_PORT} " f"Channel {CHANNEL_NAMESPACE or 'default'}" ) if os.name != "nt": return try: ctypes.windll.kernel32.SetConsoleTitleW(title) except Exception: pass class DispatcherService: def __init__(self, redis_host=REDIS_HOST, redis_port=REDIS_PORT, redis_db=REDIS_DB): self.dispatcher = self._wait_for_redis(redis_host, redis_port, redis_db) self.redis = self.dispatcher.redis version = os.environ.get("AUTOOL_DISPATCHER_VERSION", INSTANCE_NAME) self.dispatcher.monitor.start_controller_session(version=version) self.running = True self.pubsub_list = [] self._threads = [] self._lock = threading.Lock() self._report_executor = ThreadPoolExecutor(max_workers=32, thread_name_prefix="worker-report") def _wait_for_redis(self, redis_host, redis_port, redis_db, retry_interval=5): while True: try: dispatcher = RedisTaskDispatcher(redis_host=redis_host, redis_port=redis_port, redis_db=redis_db) dispatcher.redis.ping() logger.info(f"[Redis] 成功连接到 {redis_host}:{redis_port}/{redis_db}") return dispatcher except Exception as e: logger.warning(f"[Redis] 连接失败: {e}，{retry_interval}秒后重试...") time.sleep(retry_interval) def listen_worker_init(self): pubsub = self.redis.pubsub() with self._lock: self.pubsub_list.append(pubsub) pubsub.subscribe(channel_name("worker:init")) logger.info("[监听] Worker 初始化频道已启动") try: while self.running: message = pubsub.get_message(timeout=1) if message is None or message["type"] != "message": continue try: request = json.loads(message["data"]) task = self.dispatcher.worker_init( request["worker_id"], request["ip_address"], request["mac_address"], request["hostname"], request["platform"], request.get("device_type", ""), ) self.redis.publish( channel_name(f"worker:init:response:{request['worker_id']}"), json.dumps({"task": task}), ) except Exception as e: logger.error(f"处理 Worker 初始化请求时出错: {e}") finally: pubsub.unsubscribe() pubsub.close() def listen_worker_report(self): pubsub = self.redis.pubsub() with self._lock: self.pubsub_list.append(pubsub) pubsub.subscribe(channel_name("worker:report")) logger.info("[监听] Worker 上报频道已启动") try: while self.running: message = pubsub.get_message(timeout=1) if message is None or message["type"] != "message": continue try: request = json.loads(message["data"]) self._report_executor.submit(self._handle_worker_report_request, request) except Exception as e: logger.error(f"处理 Worker 上报请求时出错: {e}") finally: pubsub.unsubscribe() pubsub.close() def _handle_worker_report_request(self, request): worker_id = str(request.get("worker_id", "")).strip() if not worker_id: logger.error("处理 Worker 上报请求时出错: missing worker_id") return response_channel = channel_name(f"worker:report:response:{worker_id}") try: next_task = self.dispatcher.worker_report( worker_id, request["previous_task_key"], request.get("report_data", {}), ) payload = {"task": next_task} except Exception as e: logger.exception(f"处理 Worker 上报请求时出错: worker={worker_id}") payload = {"task": None, "error": str(e)} self.redis.publish(response_channel, json.dumps(payload)) def listen_worker_event(self): pubsub = self.redis.pubsub() with self._lock: self.pubsub_list.append(pubsub) pubsub.subscribe(channel_name("worker:event")) logger.info("[监听] Worker 事件频道已启动") try: while self.running: message = pubsub.get_message(timeout=1) if message is None or message["type"] != "message": continue try: request = json.loads(message["data"]) if str(request.get("event_type") or "").strip() == "artifacts_synced": self.dispatcher.handle_analytics_worker_event(request) else: self.dispatcher.monitor.handle_worker_event(request) except Exception as e: logger.error(f"处理 Worker 事件时出错: {e}") finally: pubsub.unsubscribe() pubsub.close() def listen_worker_retry(self): pubsub = self.redis.pubsub() with self._lock: self.pubsub_list.append(pubsub) pubsub.subscribe(channel_name("worker:retry")) logger.info("[监听] Worker 重试频道已启动") try: while self.running: message = pubsub.get_message(timeout=1) if message is None or message["type"] != "message": continue try: request = json.loads(message["data"]) success = self.dispatcher.worker_retry( request["worker_id"], request["current_task_key"], ) self.redis.publish( channel_name(f"worker:retry:response:{request['worker_id']}"), json.dumps({"success": success}), ) except Exception as e: logger.error(f"处理 Worker 重试请求时出错: {e}") finally: pubsub.unsubscribe() pubsub.close() def listen_catalog_update(self): pubsub = self.redis.pubsub() with self._lock: self.pubsub_list.append(pubsub) pubsub.subscribe(channel_name("catalog:updated")) logger.info("[监听] 应用目录更新频道已启动") try: while self.running: message = pubsub.get_message(timeout=1) if message is None or message["type"] != "message": continue try: data = json.loads(message["data"]) added = int(data.get("added") or 0) priority = str(data.get("priority") or "") if priority == "high_block": logger.info(f"[目录同步] 收到 block 任务加载通知 (added={added})，任务已直接入队，跳过 DB 刷新") continue logger.info(f"[目录同步] 收到 catalog:updated 通知 (added={added})，刷新任务队列...") new_count = self.dispatcher.refresh_tasks_from_app_summary() if new_count > 0: logger.info(f"[目录同步] 已从数据库加载 {new_count} 个新任务") else: logger.info("[目录同步] 未发现新任务") except Exception as e: logger.error(f"[目录同步] 处理 catalog:updated 消息时出错: {e}") finally: pubsub.unsubscribe() pubsub.close() def start_listeners(self): threads = [] init_thread = threading.Thread(target=self.listen_worker_init, name="worker-init-listener") init_thread.start() threads.append(init_thread) report_thread = threading.Thread(target=self.listen_worker_report, name="worker-report-listener") report_thread.start() threads.append(report_thread) event_thread = threading.Thread(target=self.listen_worker_event, name="worker-event-listener") event_thread.start() threads.append(event_thread) retry_thread = threading.Thread(target=self.listen_worker_retry, name="worker-retry-listener") retry_thread.start() threads.append(retry_thread) catalog_thread = threading.Thread(target=self.listen_catalog_update, name="catalog-update-listener") catalog_thread.start() threads.append(catalog_thread) self._threads = threads return threads def stop(self, clear_data=True): logger.info("[停止] 正在停止服务...") self.running = False self.dispatcher.monitor.close_controller_session( stop_reason="clear_data" if clear_data else "completed" ) for thread in self._threads: thread.join(timeout=3) with self._lock: for pubsub in self.pubsub_list: try: pubsub.unsubscribe() pubsub.close() except Exception: pass self.pubsub_list.clear() self._report_executor.shutdown(wait=False, cancel_futures=True) self.dispatcher.close() if clear_data: logger.info("[清理] 正在清空任务列表和 Worker 列表...") self._clear_all_data() logger.info("[清理] 清理完成") def _clear_all_data(self): try: self.redis.delete("task:queue:high") self.redis.delete("task:queue:default") self.redis.delete("task:queue:low") self.redis.delete("task:queue") self.redis.delete("task:status") self.redis.delete("task:details") self.redis.delete("task:completed") self.redis.delete("task:failed") self.redis.delete("workers:info") self.redis.delete("worker:tasks") self.redis.delete("workers:idle") self.redis.delete("workers:busy") self.redis.delete("workers:meta") for key in self.redis.scan_iter(match="alert:*"): self.redis.delete(key) for key in self.redis.scan_iter(match="cooldown:*"): self.redis.delete(key) for key in self.redis.scan_iter(match="worker:recovery_observe:*"): self.redis.delete(key) for key in self.redis.scan_iter(match="worker:recovery_pending:*"): self.redis.delete(key) logger.info("[清理] 已清空: task:queue:high/default/low, task:status, task:details, task:completed, task:failed") logger.info("[清理] 已清空: workers:info, worker:tasks, workers:idle, workers:busy, workers:meta") except Exception as e: logger.error(f"[清理] 清理数据时出错: {e}") def start_dashboard_thread(dispatcher): def _target(): try: logger.info(f"[Dashboard:{INSTANCE_NAME}] 已启动: http://{DASHBOARD_HOST}:{DASHBOARD_PORT}") run_dashboard(dispatcher=dispatcher, host=DASHBOARD_HOST, port=DASHBOARD_PORT) except Exception as e: logger.warning(f"[Dashboard] 未启动: {e}") thread = threading.Thread(target=_target, name="dashboard-server", daemon=True) thread.start() return thread def generate_final_report(dispatcher): try: os.makedirs(os.path.dirname(FINAL_REPORT_CSV) or ".", exist_ok=True) with open(FINAL_REPORT_CSV, mode="w", newline="", encoding="utf-8") as handle: writer = csv.writer(handle) writer.writerow(["任务键", "应用名称", "包名", "状态", "Worker ID", "IP地址", "开始时间", "结束时间"]) all_status = dispatcher.redis.hgetall("task:status") for task_key, status_json in sorted(all_status.items()): status = json.loads(status_json) worker_id = status.get("worker_id", "N/A") task_details_json = dispatcher.redis.hget("task:details", task_key) task_details = json.loads(task_details_json) if task_details_json else {} ip_address = "N/A" if worker_id != "N/A": worker_info_json = dispatcher.redis.hget("workers:info", worker_id) if worker_info_json: worker_info = json.loads(worker_info_json) ip_address = worker_info.get("ip_address", "N/A") writer.writerow([ task_key, task_details.get("app_name", "N/A"), task_details.get("package_name", "N/A"), status.get("status", "unknown"), worker_id, ip_address, time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(status.get("start_time", 0))), time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(status.get("end_time", 0))), ]) logger.info(f"最终报告已保存到: {FINAL_REPORT_CSV}") except Exception as e: logger.error(f"生成最终报告时出错: {e}") def generate_block_test_report(dispatcher, report_dir=None): try: output_dir = report_dir or REPORT_DIR block_report_csv = os.path.join(output_dir, "block_test_report.csv") os.makedirs(os.path.dirname(block_report_csv) or ".", exist_ok=True) with open(block_report_csv, mode="w", newline="", encoding="utf-8") as handle: writer = csv.writer(handle) writer.writerow([ "任务键", "应用名称", "包名", "状态", "Worker ID", "IP地址", "开始时间", "结束时间", "失败原因", "失败类型", "重试次数", ]) all_details = dispatcher.redis.hgetall("task:details") all_status = dispatcher.redis.hgetall("task:status") success_count = 0 failed_count = 0 retry_count = 0 for task_key in sorted(all_details.keys()): if not task_key.endswith("_block"): continue task_details_json = all_details.get(task_key) task_details = json.loads(task_details_json) if task_details_json else {} if not task_details.get("is_block_task"): continue status_json = all_status.get(task_key) status = json.loads(status_json) if status_json else {} task_status = status.get("status", "unknown") worker_id = status.get("worker_id", "N/A") ip_address = "N/A" if worker_id != "N/A": worker_info_json = dispatcher.redis.hget("workers:info", worker_id) if worker_info_json: worker_info = json.loads(worker_info_json) ip_address = worker_info.get("ip_address", "N/A") start_time = status.get("start_time", 0) end_time = status.get("end_time", 0) start_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(start_time)) if start_time else "" end_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(end_time)) if end_time else "" writer.writerow([ task_key, task_details.get("app_name", "N/A"), task_details.get("package_name", "N/A"), task_status, worker_id, ip_address, start_str, end_str, status.get("last_fail_message", ""), status.get("last_fail_type", ""), status.get("retry_count", 0), ]) if task_status in ("success", "completed", "qualified"): success_count += 1 elif task_status in ("failed",): failed_count += 1 else: retry_count += 1 writer.writerow([]) writer.writerow(["统计"]) writer.writerow(["成功", str(success_count)]) writer.writerow(["失败", str(failed_count)]) writer.writerow(["重试中/待分发", str(retry_count)]) logger.info(f"[Block测试] Block 测试报告已保存到: {block_report_csv}") except Exception as e: logger.error(f"生成 Block 测试报告时出错: {e}") def generate_worker_report(dispatcher): try: all_workers = dispatcher.get_registered_workers() os.makedirs(os.path.dirname(WORKER_REPORT_CSV) or ".", exist_ok=True) with open(WORKER_REPORT_CSV, mode="w", newline="", encoding="utf-8") as handle: writer = csv.writer(handle) writer.writerow(["Worker ID", "IP地址", "MAC地址", "主机名", "平台", "状态", "当前任务", "注册时间", "最后更新时间"]) for worker in all_workers: writer.writerow([ worker.get("worker_id", "N/A"), worker.get("ip_address", "N/A"), worker.get("mac_address", "N/A"), worker.get("hostname", "N/A"), worker.get("platform", "N/A"), worker.get("status", "N/A"), worker.get("current_task", "N/A"), time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(worker.get("register_time", 0))), time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(worker.get("last_update", 0))), ]) logger.info(f"Worker 使用报告已保存到: {WORKER_REPORT_CSV}") except Exception as e: logger.error(f"生成 Worker 报告时出错: {e}") def push_worker_status(dispatcher): try: workers = dispatcher.get_dashboard_workers() summary = dispatcher.get_dashboard_summary() online_workers = [worker for worker in workers if worker.get("online")] idle_workers = [worker for worker in online_workers if worker.get("status") == "idle"] busy_workers = [worker for worker in online_workers if worker.get("status") == "busy"] message_lines = [ "Worker 状态报告", f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", "", f"任务统计: 待分发 {summary['pending']} | 运行中 {summary['running']} | 成功 {summary['completed']} | 失败 {summary['failed']}", f"Worker 统计: 在线 {len(online_workers)} | 空闲 {len(idle_workers)} | 忙碌 {len(busy_workers)}", ] if online_workers: message_lines.append("") message_lines.append("在线 Worker:") for worker in online_workers: state = "idle" if worker.get("status") == "idle" else "busy" message_lines.append(f"- {worker.get('ip_address', 'N/A')} ({worker.get('hostname', 'N/A')}) [{state}]") dispatcher.notifier.send_weCom_alert("SYSTEM", "\n".join(message_lines)) logger.info(f"[状态推送] 已推送 Worker 状态，在线 {len(online_workers)} 台") except Exception as e: logger.error(f"[状态推送] 推送 Worker 状态时出错: {e}") def main(): _set_console_title() logger.info("=" * 60) logger.info(f"Redis 任务分发器 [{INSTANCE_NAME}]") logger.info(f"分发服务: {REDIS_HOST}:{REDIS_PORT}/{REDIS_DB}") logger.info(f"频道命名空间: {CHANNEL_NAMESPACE or 'default'}") logger.info(f"Dashboard: http://{DASHBOARD_HOST}:{DASHBOARD_PORT}") logger.info("=" * 60) service = DispatcherService(redis_host=REDIS_HOST, redis_port=REDIS_PORT, redis_db=REDIS_DB) start_dashboard_thread(service.dispatcher) logger.info("[1] 同步 Worker 注册表...") service.dispatcher.sync_worker_registry() logger.info("[2] 从应用全表加载待采集任务...") task_count = service.dispatcher.load_tasks_from_app_summary() if task_count == 0: logger.warning("应用全表中没有待采集任务，请先运行应用名单同步脚本后再启动中控") return logger.info("[3] 启动监听服务...") logger.info("等待 Worker 连接...") logger.info("按 Ctrl+C 停止程序") service.start_listeners() # 启动后立即同步一次下载队列 + 轮询一次APK结果 initial_queue_count = service.dispatcher.sync_download_queue() if initial_queue_count > 0: logger.info(f"[APK队列] 启动时已推送 {initial_queue_count} 个任务到下载队列") else: logger.info("[APK队列] 启动同步：当前无需下载的新APK任务") initial_poll_count = service.dispatcher.poll_apk_results() if initial_poll_count > 0: logger.info(f"[APK轮询] 启动轮询：导入 {initial_poll_count} 个APK") else: logger.info("[APK轮询] 启动轮询：无新APK结果") last_status_time = time.time() last_worker_status_time = time.time() last_cleanup_time = time.time() last_recovery_check_time = time.time() last_worker_push_time = time.time() last_monitor_heartbeat_time = time.time() last_task_refresh_time = time.time() last_apk_queue_time = time.time() last_apk_poll_time = time.time() last_apk_cleanup_time = time.time() last_block_report_time = time.time() tasks_done_logged = False try: while True: current_time = time.time() if current_time - last_cleanup_time > 60: cleaned_count = service.dispatcher.cleanup_stale_workers(timeout=WORKER_STALE_TIMEOUT) if cleaned_count > 0: logger.info(f"[清理] 已清理 {cleaned_count} 个超时 Worker") last_cleanup_time = current_time if current_time - last_recovery_check_time > 60: reset_workers = service.dispatcher.check_recovery_windows() if reset_workers: logger.info(f"[恢复窗口] 已重置自动恢复标记: {reset_workers}") last_recovery_check_time = current_time if current_time - last_worker_push_time >= WORKER_STATUS_PUSH_INTERVAL: push_worker_status(service.dispatcher) last_worker_push_time = current_time if current_time - last_task_refresh_time > 120: new_count = service.dispatcher.refresh_tasks_from_app_summary() if new_count > 0: logger.info(f"[任务刷新] 定期检查：从数据库加载了 {new_count} 个新任务") last_task_refresh_time = current_time if current_time - last_block_report_time >= 300: generate_block_test_report(service.dispatcher) last_block_report_time = current_time if current_time - last_apk_queue_time >= APK_DOWNLOAD_QUEUE_INTERVAL: queue_count = service.dispatcher.sync_download_queue() if queue_count > 0: logger.info(f"[APK队列] 已推送 {queue_count} 个任务到下载队列") last_apk_queue_time = current_time if current_time - last_apk_poll_time >= APK_PREFETCH_POLL_INTERVAL: imported = service.dispatcher.poll_apk_results() if imported > 0: logger.info(f"[APK轮询] 本轮导入 {imported} 个APK文件") last_apk_poll_time = current_time if current_time - last_apk_cleanup_time >= 3600: service.dispatcher._enforce_apk_storage_limits() service.dispatcher._manage_stale_apks() last_apk_cleanup_time = current_time # 心跳间隔从15秒增大到30秒，减少SQLite写入频率 if current_time - last_monitor_heartbeat_time >= 30: service.dispatcher.monitor.heartbeat_controller_session() last_monitor_heartbeat_time = current_time # 状态日志间隔从15秒增大到60秒，减少日志I/O和Redis全量查询频率 if current_time - last_status_time > 60: stats = service.dispatcher.get_statistics() logger.info( f"[状态] 待分发: {stats['pending']} | 运行中: {stats['running']} | " f"已完成: {stats['completed']} | 失败: {stats['failed']}" ) running_tasks = service.redis.hgetall("worker:tasks") if running_tasks: logger.info("[运行中任务]") for worker_id, task_key in sorted(running_tasks.items()): worker_info_json = service.redis.hget("workers:info", worker_id) if worker_info_json: worker_info = json.loads(worker_info_json) logger.info(f" - Worker {worker_id} ({worker_info.get('ip_address', 'N/A')}): {task_key}") last_status_time = current_time # Worker状态日志间隔从30秒增大到120秒，减少日志I/O if current_time - last_worker_status_time > 120: all_workers = service.dispatcher.get_dashboard_workers() idle_workers = [worker for worker in all_workers if worker.get("online") and worker.get("status") == "idle"] busy_workers = [worker for worker in all_workers if worker.get("online") and worker.get("status") == "busy"] logger.info("[Worker 状态]") logger.info(f" 总数: {len(all_workers)} | 空闲: {len(idle_workers)} | 忙碌: {len(busy_workers)}") if idle_workers: logger.info(" 空闲 Worker:") for worker in idle_workers: logger.info(f" - {worker.get('worker_id', 'N/A')} ({worker.get('ip_address', 'N/A')})") if busy_workers: logger.info(" 忙碌 Worker:") for worker in busy_workers: logger.info( f" - {worker.get('worker_id', 'N/A')} ({worker.get('ip_address', 'N/A')}) " f"-> {worker.get('current_task', 'N/A')}" ) last_worker_status_time = current_time # 使用轻量级Redis命令检测任务是否全部完成，避免每次循环都调用全量get_statistics() queue_len = service.dispatcher.get_pending_queue_length() running_count = service.redis.hlen("worker:tasks") if queue_len == 0 and running_count == 0: stats = service.dispatcher.get_statistics() if stats["pending"] == 0 and stats["running"] == 0: if not tasks_done_logged: logger.info("=" * 60) logger.info("所有任务已完成，静默等待新任务...") logger.info(f"总计: 成功 {stats['completed']} | 失败 {stats['failed']}") logger.info("=" * 60) tasks_done_logged = True else: tasks_done_logged = False else: tasks_done_logged = False # 循环间隔从2秒增大到5秒，减少轮询频率 time.sleep(5) except KeyboardInterrupt: logger.info("检测到中断信号，正在清理数据并退出...") service.stop(clear_data=True) logger.info("程序已中断退出") return except Exception as e: logger.error(f"程序发生错误: {e}", exc_info=True) service.stop(clear_data=True) return service.stop(clear_data=False) logger.info("[4] 生成最终报告...") generate_final_report(service.dispatcher) logger.info("[5] 生成 Worker 使用报告...") generate_worker_report(service.dispatcher) logger.info("[6] 生成 Block 测试报告...") generate_block_test_report(service.dispatcher) logger.info("程序结束") if __name__ == "__main__": main()