496 lines
20 KiB
Python
496 lines
20 KiB
Python
"""设备生命周期 + 通用 Worker 基类 + 全局状态注册表 + 心跳看门狗。
|
||
|
||
STFDevice — 单设备生命周期：acquire（直连/远程 server）→ release
|
||
BaseWorker — 通用 worker 基类：封装设备生命周期 + u2 连接 + 状态上报 + stop
|
||
+ 生命周期钩子(setup/teardown/on_error) + 心跳上报
|
||
子类只需实现 run_task(d)，可选实现钩子，不用管 adb/状态/异常
|
||
_WORKERS — 全局状态注册表，供 web_server 读取实时状态
|
||
_watchdog — 心跳看门狗线程，检测 worker 卡死
|
||
|
||
新增任务只需写一个 BaseWorker 子类 + 一个 BaseTask 子类，
|
||
完全不用碰这里的代码。
|
||
"""
|
||
import time
|
||
import threading
|
||
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout
|
||
|
||
import uiautomator2 as u2
|
||
|
||
from config import USB_ADB_HOST, USB_ADB_PORT
|
||
from core.logger import get_logger
|
||
from core import device_pool
|
||
from .adb_helper import adb_connect, _adb_remote
|
||
|
||
_log = get_logger("core.worker")
|
||
|
||
|
||
class STFError(Exception):
|
||
"""设备操作错误（历史名称保留兼容）。code: offline/conflict/..."""
|
||
|
||
def __init__(self, message, code=""):
|
||
super().__init__(message)
|
||
self.code = code
|
||
|
||
|
||
class DeviceOfflineError(STFError):
|
||
"""设备离线/不可用：任务不重试（换设备也没用）。"""
|
||
|
||
# 心跳超时阈值（秒）。worker 超过这个时间没更新心跳，判定为卡死
|
||
_HEARTBEAT_TIMEOUT = 120
|
||
# 看门狗检查间隔
|
||
_WATCHDOG_INTERVAL = 30
|
||
# u2.connect / d.info 超时（秒）。超过则放弃，避免 worker 永久 hang
|
||
_U2_CONNECT_TIMEOUT = 30
|
||
_U2_INFO_TIMEOUT = 10
|
||
|
||
|
||
class STFDevice:
|
||
"""单设备生命周期：acquire（直连）→ release（无操作）。
|
||
|
||
阶段 2：已摘除 STF occupy/release 与 remoteConnect 桥接——
|
||
单实例互斥由 TaskManager._running 保证；非 IP:port（USB 序列号）设备
|
||
通过 220 的远程 adb server（-H/-P）驱动，不再依赖 STF。
|
||
"""
|
||
|
||
def __init__(self, serial=None):
|
||
self.serial = serial
|
||
self.remote_adb_url = None # IP:5555 设备：设备地址
|
||
self.usb_server = None # USB 设备：远程 adb server 地址（host:port）
|
||
|
||
def acquire(self, serial=None):
|
||
self.serial = serial or self.serial or self._pick_free()
|
||
_log.info(f"[{self.serial}] 选定设备")
|
||
|
||
# 直连优先：serial 本身是 IP:5555（如 Tailscale 网络设备）时直接 adb connect。
|
||
# 只 connect、绝不 disconnect/kill-server（红线），不影响 STF provider 的连接。
|
||
if ":" in self.serial:
|
||
self.remote_adb_url = self.serial
|
||
_log.info(f"[{self.serial}] 直连设备: {self.remote_adb_url}")
|
||
if not adb_connect(self.remote_adb_url):
|
||
raise STFError(f"adb connect 失败: {self.remote_adb_url}", "offline")
|
||
else:
|
||
# USB 设备：插在 220 上，本机 adb 看不到，经远程 adb server 驱动。
|
||
# 无需 adb connect——设备已在远程 server 上，直接校验可见性。
|
||
self.usb_server = f"{USB_ADB_HOST}:{USB_ADB_PORT}"
|
||
_log.info(f"[{self.serial}] USB 设备，远程 adb server: {self.usb_server}")
|
||
out = _adb_remote(USB_ADB_HOST, USB_ADB_PORT, "devices") or ""
|
||
seen = {line.split()[0] for line in out.splitlines()[1:]
|
||
if len(line.split()) >= 2}
|
||
if self.serial not in seen:
|
||
raise STFError(f"USB 设备 {self.serial} 不在远程 adb server", "offline")
|
||
time.sleep(2)
|
||
|
||
def _pick_free(self):
|
||
free = device_pool.list_ready()
|
||
if not free:
|
||
raise STFError("没有可用的在线空闲设备", "offline")
|
||
return free[0]
|
||
|
||
def release(self):
|
||
# 直连模式：不 disconnect（红线，STF 共享 transport）；USB：无连接可断
|
||
pass
|
||
|
||
|
||
# ================== 全局 worker 状态注册表（供 web_server 读取） ==================
|
||
_WORKERS_LOCK = threading.Lock()
|
||
_WORKERS = {} # serial -> status dict
|
||
|
||
|
||
def get_all_worker_status():
|
||
with _WORKERS_LOCK:
|
||
return list(_WORKERS.values())
|
||
|
||
|
||
def _update_status(serial, **fields):
|
||
with _WORKERS_LOCK:
|
||
s = _WORKERS.setdefault(serial, {"serial": serial})
|
||
s.update(fields)
|
||
s["last_heartbeat"] = time.time()
|
||
|
||
|
||
def _remove_worker(serial):
|
||
with _WORKERS_LOCK:
|
||
_WORKERS.pop(serial, None)
|
||
|
||
|
||
def clear_worker_error(serial):
|
||
"""清除设备错误/失败状态，恢复为 idle。运行中/连接中返回 False 不清理。
|
||
|
||
检查与更新在同一把锁内完成，避免与任务启动写 connecting/running 的竞态。
|
||
"""
|
||
with _WORKERS_LOCK:
|
||
w = _WORKERS.get(serial)
|
||
if w is None:
|
||
return True # 本就没有状态记录，视作已清理
|
||
if w.get("status") in ("running", "connecting"):
|
||
return False
|
||
w.update(status="idle", last_error="", last_warning="", task_job="",
|
||
current_action="", attempt=0, progress={}, end_time=0,
|
||
last_heartbeat=time.time())
|
||
return True
|
||
|
||
|
||
def clear_all_worker_errors(exclude=()):
|
||
"""批量清除所有 error/failed 状态（exclude 里的跳过）。返回清除数。"""
|
||
exclude = set(exclude)
|
||
with _WORKERS_LOCK:
|
||
targets = [s for s, w in _WORKERS.items()
|
||
if s not in exclude and w.get("status") in ("error", "failed")]
|
||
cleared = 0
|
||
for s in targets:
|
||
if clear_worker_error(s):
|
||
cleared += 1
|
||
return cleared
|
||
|
||
|
||
def get_worker_heartbeat(serial):
|
||
"""返回 worker 最近一次心跳时间，没有返回 0。"""
|
||
with _WORKERS_LOCK:
|
||
return _WORKERS.get(serial, {}).get("last_heartbeat", 0)
|
||
|
||
|
||
# ================== 心跳看门狗 ==================
|
||
class _Watchdog(threading.Thread):
|
||
"""后台线程：检测 worker 心跳超时，自动标记异常并设置 stop 信号。
|
||
|
||
worker 卡死（如 u2 操作 hang 住）时，看门狗能发现并标记状态，
|
||
避免设备一直被占用却不干活。
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__(daemon=True, name="worker-watchdog")
|
||
self._stop = threading.Event()
|
||
|
||
def run(self):
|
||
while not self._stop.wait(_WATCHDOG_INTERVAL):
|
||
try:
|
||
self._check()
|
||
except Exception as e:
|
||
_log.error(f"看门狗检查异常: {e}")
|
||
|
||
def _check(self):
|
||
now = time.time()
|
||
with _WORKERS_LOCK:
|
||
stale = []
|
||
for serial, s in _WORKERS.items():
|
||
# 只检查运行中的 worker
|
||
if s.get("status") not in ("running", "connecting"):
|
||
continue
|
||
hb = s.get("last_heartbeat", 0)
|
||
if hb and now - hb > _HEARTBEAT_TIMEOUT:
|
||
stale.append(serial)
|
||
for serial in stale:
|
||
_log.error(f"[{serial}] 心跳超时 {now - get_worker_heartbeat(serial):.0f}s，标记卡死")
|
||
_update_status(serial, status="error",
|
||
last_error=f"心跳超时 {_HEARTBEAT_TIMEOUT}s，worker 可能卡死")
|
||
|
||
def stop(self):
|
||
self._stop.set()
|
||
|
||
|
||
_watchdog = None
|
||
|
||
|
||
def start_watchdog():
|
||
"""启动看门狗（web_server 启动时调一次）。"""
|
||
global _watchdog
|
||
if _watchdog is None or not _watchdog.is_alive():
|
||
_watchdog = _Watchdog()
|
||
_watchdog.start()
|
||
_log.info("心跳看门狗已启动")
|
||
|
||
|
||
def stop_watchdog():
|
||
"""停止看门狗。"""
|
||
global _watchdog
|
||
if _watchdog:
|
||
_watchdog.stop()
|
||
_watchdog = None
|
||
|
||
|
||
# ================== 通用 Worker 基类 ==================
|
||
class BaseWorker(threading.Thread):
|
||
"""通用 worker 基类。
|
||
|
||
自动处理：设备获取、u2 连接、状态上报、异常捕获、stop 信号、心跳。
|
||
子类只需实现 run_task(d) 方法，专注业务逻辑。
|
||
|
||
生命周期（基类 run() 已封装，不要重写）：
|
||
1. acquire 设备（直连 IP:5555 / USB 走远程 adb server）
|
||
2. u2.connect 拿到 d
|
||
3. 调用 setup(d) ← 子类可选钩子，做初始化
|
||
4. 调用子类 run_task(d) ← 业务逻辑
|
||
5. 调用 teardown(d) ← 子类可选钩子，做清理
|
||
6. finally: release 设备
|
||
|
||
子类可选钩子：
|
||
setup(d) — run_task 前的初始化（如启动 app、检查登录）。默认空。
|
||
teardown(d) — run_task 后的清理（无论成功失败）。默认空。
|
||
on_error(d, err) — run_task 抛异常时回调。默认记录日志。
|
||
|
||
子类可用工具方法：
|
||
self.stopped() — 判断是否被要求停止（循环里检查）
|
||
self.set_action(s) — 上报当前动作（同时刷新心跳）
|
||
self.set_progress() — 上报进度字段（同时刷新心跳）
|
||
self.heartbeat() — 手动刷新心跳（长操作中间调）
|
||
|
||
进度上报规范（通用，适配任意 app）：
|
||
self.set_progress(done=5, total=80, unit="视频")
|
||
- done: 已完成数量（int）
|
||
- total: 总数量（int）
|
||
- unit: 计数单位（str，如 "视频"/"轮次"/"条"）
|
||
- action_counts: 操作计数 dict（可选，如 {"like": 3, "comment": 1}）
|
||
- 也可传任意业务字段，前端会原样展示在"进度"列
|
||
示例：
|
||
self.set_progress(done=5, total=80, unit="视频",
|
||
action_counts={"like": 3, "comment": 1})
|
||
前端展示：进度条 + "5/80 视频" + 操作计数徽章
|
||
|
||
运行时长终止（通用，适配任意 app）：
|
||
self.max_duration — 最大运行时长（秒），0=不限时
|
||
self.is_time_up() — 判断是否已达 max_duration
|
||
子类在循环里检查 self.stopped() or self.is_time_up()
|
||
到时自动停（自然结束，不是 error），状态=done
|
||
|
||
异常分类（影响上层重试策略）：
|
||
DeviceOfflineError — 设备掉线/不可用，不重试（换设备也没用）
|
||
其他异常 — 可重试
|
||
"""
|
||
|
||
def __init__(self, serial, params=None, daemon=True):
|
||
super().__init__(daemon=daemon)
|
||
self.serial = serial
|
||
self.params = params or {}
|
||
self._stop_flag = threading.Event()
|
||
self.d = None # u2.Device，run_task 里用
|
||
# 通用进度字段（子类通过 set_progress 上报）
|
||
self._progress = {"done": 0, "total": 0, "unit": "",
|
||
"action_counts": {}}
|
||
# 通用运行时长终止（子类设 max_duration>0 启用）
|
||
self.max_duration = 0 # 秒，0=不限时
|
||
self._start_time = 0.0 # run_task 开始时间戳
|
||
|
||
def _start_timer(self):
|
||
"""子类在 run_task 开头调用，启动运行时长计时。"""
|
||
self._start_time = time.time()
|
||
if self.max_duration > 0:
|
||
# 预计结束时间（前端监控页展示），到点自动停止
|
||
_update_status(self.serial, end_time=self._start_time + self.max_duration)
|
||
|
||
def is_time_up(self):
|
||
"""是否已达最大运行时长。max_duration=0 时永远返回 False。"""
|
||
if self.max_duration <= 0 or self._start_time <= 0:
|
||
return False
|
||
return (time.time() - self._start_time) >= self.max_duration
|
||
|
||
def elapsed(self):
|
||
"""已运行时长（秒）。run_task 未启动时返回 0。"""
|
||
if self._start_time <= 0:
|
||
return 0
|
||
return int(time.time() - self._start_time)
|
||
|
||
def stop(self):
|
||
"""外部调用，请求停止（子类应在循环里检查 self.stopped()）。"""
|
||
self._stop_flag.set()
|
||
|
||
def stopped(self):
|
||
"""子类在循环里调用，判断是否被要求停止。"""
|
||
return self._stop_flag.is_set()
|
||
|
||
def set_action(self, action):
|
||
"""上报当前动作（子类用）。同时刷新心跳。"""
|
||
_update_status(self.serial, current_action=action)
|
||
|
||
def set_progress(self, **fields):
|
||
"""上报进度字段（子类用）。同时刷新心跳。
|
||
|
||
通用字段（前端会解析展示）：
|
||
done (int) — 已完成数量
|
||
total (int) — 总数量
|
||
unit (str) — 计数单位（"视频"/"轮次"/"条"）
|
||
action_counts (dict) — 操作计数 {"like": 3, "comment": 1}
|
||
|
||
也可传任意业务字段，会合并到 progress 里，前端原样展示。
|
||
|
||
示例：
|
||
self.set_progress(done=5, total=80, unit="视频",
|
||
action_counts={"like": 3, "comment": 1})
|
||
"""
|
||
self._progress.update(fields)
|
||
# 同步到全局状态字典，前端通过 get_status 读取
|
||
_update_status(self.serial, progress=dict(self._progress))
|
||
|
||
def heartbeat(self):
|
||
"""手动刷新心跳。长操作中间调，避免看门狗误判卡死。"""
|
||
_update_status(self.serial, last_heartbeat=time.time())
|
||
|
||
# ================== 子类实现/可选钩子 ==================
|
||
def run_task(self, d):
|
||
"""子类实现：拿到 u2.Device d，做具体业务。
|
||
|
||
不需要管 STF 占用/释放、异常、状态上报，基类已处理。
|
||
循环里记得检查 self.stopped() 以便响应停止请求。
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def setup(self, d):
|
||
"""run_task 前的初始化钩子。默认空，子类按需重写。
|
||
|
||
典型用途：启动 app、检查登录状态、进入主页。
|
||
抛异常会中断任务（走 on_error）。
|
||
"""
|
||
|
||
def teardown(self, d):
|
||
"""run_task 后的清理钩子（无论成功失败都调）。默认空，子类按需重写。
|
||
|
||
典型用途：退出 app、清理临时文件。
|
||
抛异常只记日志，不影响主流程。
|
||
"""
|
||
|
||
def on_error(self, d, error):
|
||
"""run_task 抛异常时的回调。默认空，子类按需重写。
|
||
|
||
典型用途：截图保存现场、上报错误详情。
|
||
"""
|
||
|
||
# ================== 基类主循环（不要重写） ==================
|
||
def run(self):
|
||
"""基类主循环：不要重写。子类实现 run_task + 可选钩子。"""
|
||
device = STFDevice(serial=self.serial)
|
||
_update_status(self.serial, status="connecting",
|
||
last_error="", remote_adb_url="", usb_server="", model="",
|
||
current_action="")
|
||
try:
|
||
device.acquire()
|
||
remote = device.remote_adb_url
|
||
usb_server = device.usb_server
|
||
# u2.connect 可能 hang 住（atx-agent 无响应），用线程池+超时保护；
|
||
# USB 设备（无 IP:5555）经远程 adb server 建 u2 连接
|
||
if usb_server:
|
||
self.d = self._u2_connect_usb_with_timeout(usb_server)
|
||
else:
|
||
self.d = self._u2_connect_with_timeout(remote)
|
||
if self.d is None:
|
||
raise STFError(f"u2.connect 超时 {_U2_CONNECT_TIMEOUT}s，atx-agent 无响应", "offline")
|
||
_log.info(f"[{self.serial}] u2 连接成功")
|
||
model = ""
|
||
try:
|
||
info = self._u2_info_with_timeout()
|
||
if info:
|
||
model = info.get("productName") or ""
|
||
except Exception:
|
||
pass
|
||
_update_status(self.serial, status="running",
|
||
remote_adb_url=remote or "", usb_server=usb_server or "",
|
||
model=model)
|
||
|
||
# setup 钩子
|
||
try:
|
||
self.setup(self.d)
|
||
except Exception as e:
|
||
_log.error(f"[{self.serial}] setup 异常: {e}")
|
||
raise
|
||
|
||
# 业务逻辑
|
||
_log.info(f"[{self.serial}] 进入 run_task，开始执行业务逻辑")
|
||
self.run_task(self.d)
|
||
_log.info(f"[{self.serial}] run_task 执行完毕")
|
||
|
||
# teardown 钩子（成功路径）
|
||
try:
|
||
self.teardown(self.d)
|
||
except Exception as e:
|
||
_log.warning(f"[{self.serial}] teardown 异常: {e}")
|
||
|
||
if not self.stopped():
|
||
_update_status(self.serial, status="done")
|
||
|
||
except DeviceOfflineError as e:
|
||
# 设备掉线，不重试
|
||
_log.error(f"[{self.serial}] 设备离线: {e}")
|
||
_update_status(self.serial, status="error", last_error=str(e))
|
||
self._safe_on_error(e)
|
||
except Exception as e:
|
||
err_str = str(e)
|
||
# 识别 Windows TCP 端口耗尽（WinError 10048）或 adb server 连接失败
|
||
# 这类错误是临时性的，重试时需要更长退避等端口释放
|
||
is_transient = ("10048" in err_str
|
||
or "connect to adb server failed" in err_str
|
||
or "Unable to connect to uiautomator2 server" in err_str)
|
||
if is_transient:
|
||
_log.error(f"[{self.serial}] ADB/u2 连接错误（可能端口耗尽）: {e}")
|
||
_update_status(self.serial, status="error",
|
||
last_error=f"[transient] {e}")
|
||
else:
|
||
_log.error(f"[{self.serial}] 异常: {e}", exc_info=True)
|
||
_update_status(self.serial, status="error", last_error=str(e))
|
||
self._safe_on_error(e)
|
||
finally:
|
||
try:
|
||
device.release()
|
||
except Exception as e:
|
||
_log.warning(f"[{self.serial}] release 异常: {e}")
|
||
# 只在仍在运行/连接中时才标记 released；
|
||
# done/error/failed 等终态保持不变，避免覆盖业务结果
|
||
with _WORKERS_LOCK:
|
||
cur = _WORKERS.get(self.serial, {}).get("status")
|
||
if cur in ("running", "connecting"):
|
||
_update_status(self.serial, status="released")
|
||
else:
|
||
pass
|
||
|
||
def _u2_connect_with_timeout(self, remote):
|
||
"""u2.connect 带超时保护，避免 atx-agent 无响应时永久 hang。"""
|
||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||
fut = pool.submit(u2.connect, remote)
|
||
try:
|
||
return fut.result(timeout=_U2_CONNECT_TIMEOUT)
|
||
except FuturesTimeout:
|
||
_log.error(f"[{self.serial}] u2.connect 超时 {_U2_CONNECT_TIMEOUT}s")
|
||
return None
|
||
|
||
def _u2_connect_usb_with_timeout(self, server):
|
||
"""USB 设备：经远程 adb server（220）建 u2 连接，带超时保护。
|
||
|
||
u2 3.7 底层 adbutils 支持远程 adb server：AdbClient(host, port).device(serial)，
|
||
再交给 u2.connect(AdbDevice)（与本地连接同语义，无需 adb connect）。
|
||
"""
|
||
host, _, port = server.partition(":")
|
||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||
def _connect():
|
||
from adbutils import AdbClient
|
||
dev = AdbClient(host=host, port=int(port)).device(self.serial)
|
||
return u2.connect(dev)
|
||
fut = pool.submit(_connect)
|
||
try:
|
||
return fut.result(timeout=_U2_CONNECT_TIMEOUT)
|
||
except FuturesTimeout:
|
||
_log.error(f"[{self.serial}] u2 远程连接超时 {_U2_CONNECT_TIMEOUT}s")
|
||
return None
|
||
except Exception as e:
|
||
_log.error(f"[{self.serial}] u2 远程连接失败: {e}")
|
||
return None
|
||
except Exception as e:
|
||
_log.error(f"[{self.serial}] u2.connect 异常: {e}")
|
||
return None
|
||
|
||
def _u2_info_with_timeout(self):
|
||
"""self.d.info 带超时保护，避免 atx-agent 无响应时永久 hang。"""
|
||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||
fut = pool.submit(lambda: self.d.info)
|
||
try:
|
||
return fut.result(timeout=_U2_INFO_TIMEOUT)
|
||
except FuturesTimeout:
|
||
_log.warning(f"[{self.serial}] d.info 超时 {_U2_INFO_TIMEOUT}s，跳过型号获取")
|
||
return None
|
||
except Exception as e:
|
||
_log.warning(f"[{self.serial}] d.info 异常: {e}")
|
||
return None
|
||
|
||
def _safe_on_error(self, error):
|
||
"""安全调用 on_error 钩子，不让钩子异常影响主流程。"""
|
||
try:
|
||
self.on_error(self.d, error)
|
||
except Exception as e:
|
||
_log.warning(f"[{self.serial}] on_error 钩子异常: {e}")
|