337 lines
12 KiB
Python
Executable File
337 lines
12 KiB
Python
Executable File
# -*- 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()
|