# -*- encoding=utf8 -*- import redis from redis import ConnectionPool import json import time import socket import uuid import platform from typing import Optional, Dict, Any, Tuple from config import REDIS_HOST, REDIS_PORT, REDIS_DB, channel_name, normalize_worker_id class TaskWorker: _connection_pool: Optional[ConnectionPool] = None _connection_pool_config: Optional[Tuple[str, int, int, int]] = None def __init__(self, redis_host=REDIS_HOST, redis_port=REDIS_PORT, redis_db=REDIS_DB, max_connections=10): pool_config = (redis_host, redis_port, redis_db, max_connections) if TaskWorker._connection_pool is None or TaskWorker._connection_pool_config != pool_config: TaskWorker._connection_pool = ConnectionPool( host=redis_host, port=redis_port, db=redis_db, decode_responses=True, max_connections=max_connections, socket_timeout=30, socket_connect_timeout=10, retry_on_timeout=True ) TaskWorker._connection_pool_config = pool_config self.redis = redis.Redis(connection_pool=TaskWorker._connection_pool) self.current_task: Optional[Dict[str, Any]] = None self.worker_id: Optional[str] = None self.ip_address: Optional[str] = None self.mac_address: Optional[str] = None self.hostname: Optional[str] = None self.platform: Optional[str] = None self.DEFAULT_TIMEOUT = 60 self.WAIT_INTERVAL = 30 def _get_ip_address(self): """获取本机IP地址（优先获取192.168开头的IP）""" try: hostname = socket.gethostname() ip_addresses = socket.gethostbyname_ex(hostname)[2] for ip in ip_addresses: if ip.startswith('192.168'): return ip if ip_addresses: return ip_addresses[0] return '127.0.0.1' except Exception as e: print(f"获取IP地址失败: {e}") return '127.0.0.1' def _get_mac_address(self) -> str: """获取本机Mac地址""" try: mac = uuid.getnode() mac_address = ':'.join(['{:02x}'.format((mac >> elements) & 0xff) for elements in range(0, 8*6, 8)][::-1]) return mac_address.upper() except Exception as e: print(f"获取MAC地址失败: {e}") return '00:00:00:00:00:00' def _send_request(self, publish_channel: str, response_channel: str, request_data: Dict[str, Any], timeout: Optional[int] = None) -> Optional[Dict[str, Any]]: """通用的请求发送方法，避免代码重复 Args: publish_channel: 发布请求的频道 response_channel: 订阅响应的频道 request_data: 请求数据 timeout: 超时时间(秒)，默认使用self.DEFAULT_TIMEOUT Returns: 响应数据字典，超时或失败返回None """ if timeout is None: timeout = self.DEFAULT_TIMEOUT pubsub = self.redis.pubsub() try: pubsub.subscribe(response_channel) self.redis.publish(publish_channel, json.dumps(request_data)) start_time = time.time() while time.time() - start_time < timeout: message = pubsub.get_message(timeout=1) if message and message['type'] == 'message': return json.loads(message['data']) return None finally: pubsub.unsubscribe(response_channel) pubsub.close() def init(self) -> Optional[Dict[str, Any]]: """初始化函数：上报MAC+IP并领取第一个任务 Returns: dict: 包含任务信息的字典 如果没有任务，会阻塞等待直到有任务 """ try: self.ip_address = self._get_ip_address() self.mac_address = self._get_mac_address() self.hostname = socket.gethostname() self.platform = platform.system() self.worker_id = normalize_worker_id(ip_address=self.ip_address) print(f"Worker信息:") print(f" Worker ID: {self.worker_id}") print(f" IP地址: {self.ip_address}") print(f" MAC地址: {self.mac_address}") print(f" 主机名: {self.hostname}") print(f" 平台: {self.platform}") init_request = { 'worker_id': self.worker_id, 'ip_address': self.ip_address, 'mac_address': self.mac_address, 'hostname': self.hostname, 'platform': self.platform } while True: print("等待分发器响应...") response = self._send_request( publish_channel=channel_name("worker:init"), response_channel=channel_name(f"worker:init:response:{self.worker_id}"), request_data=init_request ) if response is None: print(f"初始化超时，等待 {self.WAIT_INTERVAL} 秒后重试...") time.sleep(self.WAIT_INTERVAL) continue task = response.get('task') if task: self.current_task = task print(f"初始化成功，领取任务: {task['app_name']} ({task['package_name']})") return task else: print(f"暂无任务，等待 {self.WAIT_INTERVAL} 秒后重试...") time.sleep(self.WAIT_INTERVAL) except Exception as e: print(f"初始化失败: {e}") import traceback traceback.print_exc() return None def report(self, status: str, message: str = '', error_type: str = 'UNKNOWN') -> Optional[Dict[str, Any]]: """上报前一个任务的完成状态，并获取下一个任务 Args: status: 任务完成状态 - 'success': 任务成功，获取下一个任务 - 'failed': 任务失败，可重试，获取下一个任务 - 'stop': 任务失败，Worker请求停止，不再接收新任务 message: 附加消息，通常为空，stop状态时存放告警信息 error_type: 失败类型（如 TIMEOUT, CRASH, ASSERTION_ERROR 等） Returns: dict: 包含下一个任务信息的字典 如果没有任务，会阻塞等待直到有任务（stop状态除外） """ try: if not self.current_task: print("警告: 没有当前任务，无法上报") return None previous_task_key = self.current_task['task_key'] print(f"上报任务: {previous_task_key} = {status}") if message: print(f"附加消息: {message}") if error_type != 'UNKNOWN': print(f"失败类型: {error_type}") report_request = { 'worker_id': self.worker_id, 'previous_task_key': previous_task_key, 'status': status, 'message': message, 'error_type': error_type } if status == 'stop': print("发送停止请求...") response = self._send_request( publish_channel=channel_name("worker:report"), response_channel=channel_name(f"worker:report:response:{self.worker_id}"), request_data=report_request ) self.current_task = None return None while True: print("等待分发器响应...") response = self._send_request( publish_channel=channel_name("worker:report"), response_channel=channel_name(f"worker:report:response:{self.worker_id}"), request_data=report_request ) if response is None: print(f"上报超时，等待 {self.WAIT_INTERVAL} 秒后重试...") time.sleep(self.WAIT_INTERVAL) continue next_task = response.get('task') if next_task: self.current_task = next_task print(f"上报成功，领取新任务: {next_task['app_name']} ({next_task['package_name']})") return next_task else: print("当前无新任务，结束本轮上报") self.current_task = None return None except Exception as e: print(f"上报失败: {e}") import traceback traceback.print_exc() return None def retry(self) -> bool: """重试函数：告知该任务需重新计时 Returns: bool: 是否成功更新 """ try: if not self.current_task: print("警告: 没有当前任务，无法发送重试请求") return False retry_request = { 'worker_id': self.worker_id, 'current_task_key': self.current_task['task_key'] } response = self._send_request( publish_channel=channel_name("worker:retry"), response_channel=channel_name(f"worker:retry:response:{self.worker_id}"), request_data=retry_request, timeout=5 ) if response is None: print("重试请求超时") return False success = response.get('success', False) if success: print(f"重试成功，任务 {self.current_task['task_key']} 已重新计时") else: print("重试失败") return success except Exception as e: print(f"重试失败: {e}") return False def main(): print("=" * 60) print("Redis任务分发Worker（简化接口版）") print("=" * 60) worker = TaskWorker() try: print("\n[1] 初始化Worker...") task = worker.init() if not task: print("初始化失败，程序退出") return print("\n[2] 开始处理任务...") print("按 Ctrl+C 停止Worker\n") while task: try: app_name = task['app_name'] package_name = task['package_name'] task_key = task['task_key'] print(f"\n开始处理任务:") print(f" 应用名称: {app_name}") print(f" 包名: {package_name}") print(f" 任务键: {task_key}") print("正在处理任务...") time.sleep(2) print("发送心跳...") worker.retry() time.sleep(10) status = 'success' print(f"任务完成，状态: {status}") task = worker.report(status) except KeyboardInterrupt: print("\n检测到中断信号，正在退出...") break except Exception as e: print(f"处理任务时出错: {e}") import traceback traceback.print_exc() task = worker.report('failed') print("\n所有任务已完成或无更多任务") except KeyboardInterrupt: print("\n检测到中断信号，正在退出...") except Exception as e: print(f"\n程序发生错误: {e}") import traceback traceback.print_exc() if __name__ == "__main__": main()