cmautobuy/client/src/pdd_device_service.py

"""uiautomator2 设备连接边界。

普通服务每个任务建立一次连接；持久服务只复用 Device 连接，不复用页面数据。
两种会话都只能在创建它们的线程中使用。
"""

from __future__ import annotations

import threading
import time
from contextlib import AbstractContextManager, contextmanager
from contextvars import ContextVar
from typing import Any, Callable, Iterator, Optional

from .adb_runtime import BundledAdbError, configure_bundled_adb_environment
from .performance_timing import current_performance_trace


PDD_PACKAGE_NAME = "com.xunmeng.pinduoduo"


class PddDeviceError(RuntimeError):
    """设备连接失败，并携带供任务状态使用的稳定错误码。"""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


class PddDeviceSession(AbstractContextManager[Any]):
    """一个只能由创建线程使用的 uiautomator2 Device 会话。"""

    def __init__(
        self,
        device: Any,
        serial: str,
        initial_app_state: dict[str, Any],
        release: Callable[[], None],
    ) -> None:
        self._device = device
        self.serial = serial
        self.initial_app_state = dict(initial_app_state)
        self._release = release
        self._owner_thread_id = threading.get_ident()
        self._closed = False

    @property
    def device(self) -> Any:
        """返回 Device；跨线程或会话关闭后访问会明确失败。"""

        if self._closed:
            raise PddDeviceError("DEVICE_SESSION_CLOSED", "设备会话已经关闭")
        if threading.get_ident() != self._owner_thread_id:
            raise PddDeviceError(
                "DEVICE_THREAD_VIOLATION",
                "uiautomator2 Device 只能在创建会话的工作线程中使用",
            )
        return self._device

    def __enter__(self) -> Any:
        return self.device

    def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None:
        if not self._closed:
            self._closed = True
            self._release()
        return None


class PddDeviceService:
    """连接并校验一台已保存的 Android 设备。"""

    def __init__(
        self,
        connector: Optional[Callable[[str], Any]] = None,
    ) -> None:
        self._connector = connector or self._default_connector
        self._state_lock = threading.Lock()
        self._active_serial: Optional[str] = None

    @staticmethod
    def _default_connector(serial: str) -> Any:
        try:
            configure_bundled_adb_environment()
        except BundledAdbError as exc:
            raise PddDeviceError("DEVICE_ADB_MISSING", str(exc)) from exc
        try:
            import uiautomator2 as u2
        except ImportError as exc:
            raise PddDeviceError(
                "DEVICE_U2_MISSING",
                "未安装 uiautomator2，请先安装 client/requirements.txt",
            ) from exc
        return u2.connect(serial)

    @staticmethod
    def _validate_serial(serial: str) -> str:
        value = str(serial or "").strip()
        if not value or any(character.isspace() for character in value):
            raise PddDeviceError("DEVICE_ADDRESS_INVALID", "Android 设备号无效")
        return value

    def connect(self, serial: str) -> PddDeviceSession:
        """连接设备并返回单线程独占会话。

        调用方必须使用 ``with service.connect(serial) as device``。同一个服务
        已有活动会话时会拒绝第二次连接。
        """

        checked_serial = self._validate_serial(serial)
        with self._state_lock:
            if self._active_serial is not None:
                raise PddDeviceError(
                    "DEVICE_IN_USE",
                    f"设备 {self._active_serial} 正在执行其他自动化任务",
                )
            self._active_serial = checked_serial

        try:
            trace = current_performance_trace()
            if trace is None:
                device = self._connector(checked_serial)
            else:
                with trace.stage("uiautomator2_connect"):
                    device = self._connector(checked_serial)
            if trace is None:
                current = device.app_current()
            else:
                with trace.stage("first_app_current"):
                    current = device.app_current()
            if not isinstance(current, dict):
                raise RuntimeError("uiautomator2 未返回有效的设备状态")
        except PddDeviceError:
            self._release(checked_serial)
            raise
        except Exception as exc:
            self._release(checked_serial)
            details = str(exc).lower()
            code = (
                "DEVICE_OFFLINE"
                if any(word in details for word in ("offline", "not found", "disconnected"))
                else "DEVICE_CONNECT_FAILED"
            )
            raise PddDeviceError(
                code,
                f"无法连接 Android 设备 {checked_serial}：{exc}",
            ) from exc

        return PddDeviceSession(
            device,
            checked_serial,
            current,
            lambda: self._release(checked_serial),
        )

    def _release(self, serial: str) -> None:
        with self._state_lock:
            if self._active_serial == serial:
                self._active_serial = None


class PersistentPddDeviceService(PddDeviceService):
    """在固定工作线程内复用同一台设备的 Device 连接。

    ``connect`` 每次仍调用 ``app_current`` 做健康检查并建立全新的任务会话。
    缓存连接失效时最多重连一次；页面 XML、选择器和业务数据均不缓存。
    """

    def __init__(
        self,
        connector: Optional[Callable[[str], Any]] = None,
        *,
        ttl_seconds: float = 90.0,
        monotonic: Callable[[], float] = time.monotonic,
    ) -> None:
        super().__init__(connector)
        if ttl_seconds <= 0:
            raise ValueError("设备会话 TTL 必须大于 0")
        self._ttl_seconds = float(ttl_seconds)
        self._monotonic = monotonic
        self._owner_thread_id: Optional[int] = None
        self._cached_serial: Optional[str] = None
        self._cached_device: Any = None
        self._last_used_at = 0.0

    @property
    def ttl_seconds(self) -> float:
        return self._ttl_seconds

    @property
    def has_cached_device(self) -> bool:
        """只供生命周期管理和测试判断，不返回 Device 对象。"""

        return self._cached_device is not None

    def connect(self, serial: str) -> PddDeviceSession:
        """在所有者线程取得独占任务会话，并按需复用底层连接。"""

        self._assert_owner_thread()
        checked_serial = self._validate_serial(serial)
        with self._state_lock:
            if self._active_serial is not None:
                raise PddDeviceError(
                    "DEVICE_IN_USE",
                    f"设备 {self._active_serial} 正在执行其他自动化任务",
                )
            self._active_serial = checked_serial

        try:
            device, current = self._healthy_device(checked_serial)
        except PddDeviceError:
            self._release(checked_serial)
            raise
        except Exception as exc:
            self._release(checked_serial)
            details = str(exc).lower()
            code = (
                "DEVICE_OFFLINE"
                if any(
                    word in details
                    for word in ("offline", "not found", "disconnected")
                )
                else "DEVICE_CONNECT_FAILED"
            )
            raise PddDeviceError(
                code,
                f"无法连接 Android 设备 {checked_serial}：{exc}",
            ) from exc

        return PddDeviceSession(
            device,
            checked_serial,
            current,
            lambda: self._release_persistent(checked_serial),
        )

    def release_cached(self) -> None:
        """在所有者线程丢弃缓存连接；可重复调用。"""

        self._assert_owner_thread()
        with self._state_lock:
            if self._active_serial is not None:
                raise PddDeviceError(
                    "DEVICE_IN_USE",
                    "设备任务尚未到达安全结束点，暂不能释放连接",
                )
            self._discard_cached()

    def _healthy_device(self, serial: str) -> tuple[Any, dict[str, Any]]:
        now = self._monotonic()
        expired = (
            self._cached_device is not None
            and now - self._last_used_at >= self._ttl_seconds
        )
        if self._cached_serial != serial or expired:
            self._discard_cached()

        if self._cached_device is not None:
            try:
                current = self._read_app_state(self._cached_device)
                return self._cached_device, current
            except Exception:
                # 旧连接失效时只丢弃并重连一次，不能在这里无限重试。
                self._discard_cached()

        trace = current_performance_trace()
        if trace is None:
            device = self._connector(serial)
        else:
            with trace.stage("uiautomator2_connect"):
                device = self._connector(serial)
        current = self._read_app_state(device)
        self._cached_serial = serial
        self._cached_device = device
        self._last_used_at = now
        return device, current

    @staticmethod
    def _read_app_state_without_trace(device: Any) -> dict[str, Any]:
        current = device.app_current()
        if not isinstance(current, dict):
            raise RuntimeError("uiautomator2 未返回有效的设备状态")
        return current

    def _read_app_state(self, device: Any) -> dict[str, Any]:
        trace = current_performance_trace()
        if trace is None:
            return self._read_app_state_without_trace(device)
        with trace.stage("first_app_current"):
            return self._read_app_state_without_trace(device)

    def _release_persistent(self, serial: str) -> None:
        self._assert_owner_thread()
        with self._state_lock:
            if self._active_serial == serial:
                self._active_serial = None
                self._last_used_at = self._monotonic()

    def _discard_cached(self) -> None:
        # uiautomator2 Device 没有稳定的 close API；清除唯一引用即可让底层
        # HTTP 客户端按库自身生命周期回收，不能猜测调用私有方法。
        self._cached_serial = None
        self._cached_device = None
        self._last_used_at = 0.0

    def _assert_owner_thread(self) -> None:
        current = threading.get_ident()
        if self._owner_thread_id is None:
            self._owner_thread_id = current
        elif self._owner_thread_id != current:
            raise PddDeviceError(
                "DEVICE_THREAD_VIOLATION",
                "持久 uiautomator2 Device 只能在固定工作线程中使用和释放",
            )


_CURRENT_DEVICE_SERVICE: ContextVar[Optional[PersistentPddDeviceService]] = (
    ContextVar("pdd_device_service", default=None)
)


@contextmanager
def bind_thread_device_service(
    service: PersistentPddDeviceService,
) -> Iterator[None]:
    """让当前工作线程创建的采集、采购和核单适配器使用同一连接服务。"""

    token = _CURRENT_DEVICE_SERVICE.set(service)
    try:
        yield
    finally:
        _CURRENT_DEVICE_SERVICE.reset(token)


def current_thread_device_service() -> Optional[PersistentPddDeviceService]:
    """返回当前工作线程绑定的服务；未绑定时返回 ``None``。"""

    return _CURRENT_DEVICE_SERVICE.get()