509 lines
19 KiB
Python
509 lines
19 KiB
Python
# -*- encoding=utf8 -*-
|
||
import redis
|
||
from redis import ConnectionPool
|
||
import json
|
||
import time
|
||
import socket
|
||
import uuid
|
||
import platform
|
||
import os
|
||
import sys
|
||
from typing import Optional, Dict, Any, Tuple
|
||
|
||
CONFIG_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
if CONFIG_ROOT not in sys.path:
|
||
sys.path.insert(0, CONFIG_ROOT)
|
||
|
||
from config_loader import load_config as load_autool_config
|
||
|
||
def _load_worker_config() -> Dict[str, Any]:
|
||
return load_autool_config()
|
||
|
||
|
||
def _resolve_device_type(payload: Dict[str, Any]) -> str:
|
||
return "emulator" if bool(payload.get("IS_EMULATOR", True)) else "physical"
|
||
|
||
|
||
WORKER_CONFIG = _load_worker_config()
|
||
CONTROL_REDIS_HOST = str(WORKER_CONFIG["CONTROL_REDIS_HOST"]).strip()
|
||
CONTROL_REDIS_PORT = int(WORKER_CONFIG["CONTROL_REDIS_PORT"])
|
||
CONTROL_REDIS_DB = int(WORKER_CONFIG["CONTROL_REDIS_DB"])
|
||
CONTROL_CHANNEL_NAMESPACE = str(WORKER_CONFIG.get("CONTROL_CHANNEL_NAMESPACE", "")).strip()
|
||
DEFAULT_DEVICE_TYPE = _resolve_device_type(WORKER_CONFIG)
|
||
|
||
|
||
def _channel_name(name: str) -> str:
|
||
return f"{CONTROL_CHANNEL_NAMESPACE}:{name}" if CONTROL_CHANNEL_NAMESPACE else name
|
||
|
||
|
||
class TaskWorker:
|
||
_connection_pool: Optional[ConnectionPool] = None
|
||
_connection_pool_config: Optional[Tuple[str, int, int, int]] = None
|
||
|
||
def __init__(
|
||
self,
|
||
redis_host=CONTROL_REDIS_HOST,
|
||
redis_port=CONTROL_REDIS_PORT,
|
||
redis_db=CONTROL_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.device_type: Optional[str] = None
|
||
self.DEFAULT_TIMEOUT = 60
|
||
self.WAIT_INTERVAL = 30
|
||
|
||
def _get_ip_address(self):
|
||
"""获取本机IP地址（优先获取192.168.2.开头的IP）"""
|
||
try:
|
||
hostname = socket.gethostname()
|
||
ip_addresses = socket.gethostbyname_ex(hostname)[2]
|
||
|
||
for ip in ip_addresses:
|
||
if ip.startswith('192.168.2.'):
|
||
return ip
|
||
|
||
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地址，优先匹配真实物理网卡，保持与add_workers.py逻辑一致"""
|
||
if platform.system() == "Windows":
|
||
try:
|
||
import subprocess
|
||
import re
|
||
result = subprocess.run(["getmac", "/fo", "csv", "/nh"], capture_output=True, text=True, timeout=10)
|
||
output = result.stdout.strip()
|
||
|
||
best_mac = None
|
||
for line in output.splitlines():
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
parts = [p.strip().strip('"') for p in line.split(",")]
|
||
if len(parts) < 2:
|
||
continue
|
||
mac = parts[0].strip()
|
||
transport = parts[1].strip() if len(parts) > 1 else ""
|
||
|
||
# 跳过无效 MAC
|
||
if not re.match(r"([0-9A-Fa-f]{2}[:-]){5}[0-9A-Fa-f]{2}", mac):
|
||
continue
|
||
|
||
# 跳过已断开的连接
|
||
if "已断开" in transport or "Disconnected" in transport.lower() or "Media disconnected" in transport.lower():
|
||
continue
|
||
|
||
mac = mac.upper().replace("-", ":")
|
||
|
||
# 优先选取硬件以太网适配器
|
||
transport_lower = transport.lower()
|
||
if any(kw in transport_lower for kw in ["ethernet", "以太网", "realtek", "intel"]):
|
||
best_mac = mac
|
||
break
|
||
|
||
if best_mac is None:
|
||
best_mac = mac
|
||
|
||
if best_mac:
|
||
return best_mac
|
||
except Exception as e:
|
||
print(f"通过getmac获取MAC地址失败，尝试降级方案: {e}")
|
||
|
||
# 降级方案（非Windows或getmac失败）
|
||
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.device_type = DEFAULT_DEVICE_TYPE
|
||
self.worker_id = f"{self.ip_address}_{self.mac_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}")
|
||
print(f" 设备类型: {self.device_type}")
|
||
|
||
init_request = {
|
||
'worker_id': self.worker_id,
|
||
'ip_address': self.ip_address,
|
||
'mac_address': self.mac_address,
|
||
'hostname': self.hostname,
|
||
'platform': self.platform,
|
||
'device_type': self.device_type,
|
||
}
|
||
|
||
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("暂无任务，等待 {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, report_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
||
"""上报前一个任务的完成状态，并获取下一个任务
|
||
|
||
Args:
|
||
report_data: 上报数据字典，结构如下:
|
||
{
|
||
"status": "success|failed|stop", # 必填
|
||
"error": { # 失败时必填
|
||
"type": "DOWNLOAD_ERROR/5",
|
||
"category": "DOWNLOAD_ERROR",
|
||
"code": 5,
|
||
"reason": "错误原因描述",
|
||
"details": "详细错误信息（可选）",
|
||
"crashed_source": "google_play" # APP_CRASH时必填，闪退的下载源
|
||
},
|
||
"metrics": { # 可选
|
||
"login_count": 0,
|
||
"register_count": 0,
|
||
"guiagent_message": "GuiAgent消息",
|
||
"scenario_triggered": true
|
||
},
|
||
"statistics": { # 可选
|
||
"round": 1,
|
||
"exit_code": 0,
|
||
"error_reason": "...",
|
||
"droidbot_steps": 100,
|
||
"guiagent_steps": 50,
|
||
"total_steps": 150,
|
||
"duration_seconds": 120.5,
|
||
"is_retry": false,
|
||
"download_source": "google"
|
||
}
|
||
}
|
||
|
||
Returns:
|
||
dict: 包含下一个任务信息的字典
|
||
如果没有任务，会阻塞等待直到有任务（stop状态除外）
|
||
"""
|
||
try:
|
||
if not self.current_task:
|
||
print("警告: 没有当前任务，无法上报")
|
||
return None
|
||
|
||
previous_task_key = self.current_task['task_key']
|
||
status = report_data.get('status', 'failed')
|
||
|
||
print(f"上报任务: {previous_task_key} = {status}")
|
||
|
||
error_info = report_data.get('error', {})
|
||
if error_info:
|
||
error_type = error_info.get('type')
|
||
if not error_type:
|
||
category = error_info.get('category')
|
||
code = error_info.get('code')
|
||
error_type = f"{category}/{code}" if category is not None and code is not None else 'UNKNOWN'
|
||
error_reason = error_info.get('reason', '')
|
||
print(f"错误类型: {error_type}")
|
||
if error_reason:
|
||
print(f"错误原因: {error_reason}")
|
||
|
||
metrics = report_data.get('metrics', {})
|
||
if metrics:
|
||
login_count = metrics.get('login_count', 0)
|
||
register_count = metrics.get('register_count', 0)
|
||
if login_count > 0 or register_count > 0:
|
||
print(f"指标: 登录场景 {login_count} 次, 注册场景 {register_count} 次")
|
||
|
||
statistics = report_data.get('statistics', {})
|
||
if statistics:
|
||
duration = statistics.get('duration_seconds', 0)
|
||
total_steps = statistics.get('total_steps', 0)
|
||
print(f"统计: 耗时 {duration}s, 总步数 {total_steps}")
|
||
|
||
report_request = {
|
||
'worker_id': self.worker_id,
|
||
'previous_task_key': previous_task_key,
|
||
'report_data': report_data
|
||
}
|
||
|
||
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 event(self, event_data: Dict[str, Any]) -> bool:
|
||
"""发送监控事件，不阻塞主流程。"""
|
||
try:
|
||
if not self.worker_id:
|
||
return False
|
||
payload = dict(event_data or {})
|
||
payload.setdefault("worker_id", self.worker_id)
|
||
payload.setdefault("event_time", time.time())
|
||
if "task_key" not in payload and self.current_task:
|
||
payload["task_key"] = self.current_task.get("task_key")
|
||
self.redis.publish(_channel_name("worker:event"), json.dumps(payload))
|
||
return True
|
||
except Exception as e:
|
||
print(f"发送监控事件失败: {e}")
|
||
return False
|
||
|
||
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']
|
||
}
|
||
|
||
# 使用通用方法发送请求（5秒超时）
|
||
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:
|
||
# 1. 初始化：注册Worker并领取第一个任务
|
||
print("\n[1] 初始化Worker...")
|
||
task = worker.init()
|
||
|
||
if not task:
|
||
print("没有任务需要处理，程序退出")
|
||
return
|
||
|
||
# 2. 处理任务循环
|
||
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' # 或 'failed'
|
||
print(f"任务完成，状态: {status}")
|
||
|
||
# 3. 上报任务状态并获取下一个任务
|
||
report_data = {
|
||
"status": status,
|
||
"error": None,
|
||
"metrics": {
|
||
"login_count": 0,
|
||
"register_count": 0
|
||
}
|
||
}
|
||
task = worker.report(report_data)
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n检测到中断信号，正在退出...")
|
||
break
|
||
except Exception as e:
|
||
print(f"处理任务时出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# 上报失败状态
|
||
task = worker.report({
|
||
"status": "failed",
|
||
"error": {"type": "TASK_ERROR", "reason": str(e)}
|
||
})
|
||
|
||
print("\n所有任务已完成或无更多任务")
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n检测到中断信号，正在退出...")
|
||
except Exception as e:
|
||
print(f"\n程序发生错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
if __name__ == "__main__":
|
||
main()
|