kefu/wechat_rpa/send_lock.py

"""
发送互斥锁
==========
双引擎并存（引擎 A 截图 RPA + 引擎 B 数据直读）时，企业微信窗口是唯一共享
资源，任意时刻只允许一个发送者触碰它。本模块用 Windows 命名互斥量实现
跨进程/跨线程的发送互斥：

  - 拿到锁的引擎才有资格调用 WeChatBot.send_reply()；
  - 拿不到的引擎本轮直接跳过，绝不阻塞等待（等待会导致两个引擎互相拖死）；
  - 持锁进程/线程崩溃时，Windows 会自动释放互斥量，不会出现"锁泄漏"。

非 Windows 环境（CI、单元测试、macOS/Linux 开发机）回退到进程内
threading.Lock，保证模块可导入、可测试。
"""

import os
import threading
import time

# 全局唯一的互斥量名。跨进程共享同一命名空间，任意数量的引擎进程都抢同一把锁。
_MUTEX_NAME = "Global\\WeChatBotSendLock"

# WaitForSingleObject 返回码
_WAIT_OBJECT_0 = 0x00000000
_WAIT_ABANDONED = 0x00000080   # 拿到了锁，但上一个持有者是崩溃退出的
_WAIT_TIMEOUT = 0x00000102

# 回退锁（非 Windows）：进程内互斥，行为与命名互斥量一致
_FALLBACK_LOCK = threading.Lock()

_handle = None
_init_done = False
_owner = ""          # 当前持有者标识（引擎名:线程名），用于排查与日志
_owner_until = 0.0   # 持有者预计释放时间（wall clock），仅用于日志统计
_holder_tid = None   # 持有锁的线程 id；防止同一线程重入（嵌套发送）


def _windows_mutex_available() -> bool:
    """Windows 且能加载 kernel32 时返回 True。"""
    global _handle, _init_done
    if _init_done:
        return _handle is not None
    _init_done = True
    if os.name != "nt":
        return False
    try:
        import ctypes
        from ctypes import wintypes  # noqa: F401

        _handle = ctypes.windll.kernel32.CreateMutexW(None, False, _MUTEX_NAME)
        return bool(_handle)
    except Exception:
        _handle = None
        return False


def try_acquire(owner: str = "", timeout: float = 0.0) -> bool:
    """非阻塞（或短超时）抢占发送锁。

    owner 用于日志排查，例如 "engine_b:listener"。timeout 默认 0（不等待），
    调用方应优先用 0——拿不到就下一轮再试，避免双引擎互相阻塞。

    注意：Windows 命名互斥量对同一线程是可重入的（可重复获取），这不符合
    「任意时刻只有一个发送者」的语义（同一线程内也不允许嵌套发送）。因此
    本函数额外记录持有线程 id，同线程重入一律返回 False。
    """
    global _owner, _owner_until, _holder_tid
    tid = threading.get_ident()
    if _holder_tid == tid:
        return False
    if _windows_mutex_available():
        try:
            import ctypes

            wait_ms = int(max(0.0, timeout) * 1000)
            rc = ctypes.windll.kernel32.WaitForSingleObject(_handle, wait_ms)
            # WAIT_ABANDONED 同样代表"这把锁现在归你了"，只是上一个持有者是崩溃
            # 退出的。把它当失败直接 return，互斥量就永远留在本线程名下、再也没人
            # 释放——从此整个进程一条消息都发不出去。
            if rc in (_WAIT_OBJECT_0, _WAIT_ABANDONED):
                _owner = str(owner or "")
                _owner_until = time.time() + max(0.0, timeout)
                _holder_tid = tid
                return True
            return False
        except Exception:
            return False
    # threading.Lock 不允许"非阻塞 + 超时"同时出现（ValueError）。默认的
    # timeout=0 正是这种组合，非 Windows 环境下每一次抢锁都会直接抛异常。
    if timeout > 0:
        acquired = _FALLBACK_LOCK.acquire(blocking=True, timeout=timeout)
    else:
        acquired = _FALLBACK_LOCK.acquire(blocking=False)
    if not acquired:
        return False
    _owner = str(owner or "")
    _owner_until = time.time() + max(0.0, timeout)
    _holder_tid = tid
    return True


def release() -> None:
    """释放发送锁。未持锁时调用是安全的（Windows 互斥量允许）。"""
    global _owner, _owner_until, _holder_tid
    if _holder_tid is not None and _holder_tid != threading.get_ident():
        # 非持有线程尝试释放：忽略，避免误释放他人的锁
        return
    _holder_tid = None
    if _windows_mutex_available():
        try:
            import ctypes

            ctypes.windll.kernel32.ReleaseMutex(_handle)
        except Exception:
            pass
    else:
        try:
            _FALLBACK_LOCK.release()
        except RuntimeError:
            pass
    _owner = ""
    _owner_until = 0.0


def is_held() -> bool:
    """当前进程是否持有发送锁（用于诊断日志）。"""
    if _holder_tid is not None:
        return True
    if _windows_mutex_available():
        try:
            import ctypes

            rc = ctypes.windll.kernel32.WaitForSingleObject(_handle, 0)
            if rc in (_WAIT_OBJECT_0, _WAIT_ABANDONED):
                # 抢到了说明之前没持有，立即释放还原状态。WAIT_ABANDONED 也已经
                # 把所有权交给了本线程，同样必须还回去，否则这次"诊断"会把锁吞掉。
                ctypes.windll.kernel32.ReleaseMutex(_handle)
                return False
            return rc == _WAIT_TIMEOUT  # 被别人持有
        except Exception:
            return False
    return _FALLBACK_LOCK.locked()


def holder() -> str:
    """当前持有者标识（仅本进程视角；跨进程时返回空串）。"""
    return str(_owner or "")


# ── 上下文管理器：with send_lock("engine_b"): ... ───────────────────────────
class _SendLockCtx:
    def __init__(self, owner: str, timeout: float):
        self.owner = owner
        self.timeout = timeout
        self.acquired = False

    def __enter__(self):
        self.acquired = try_acquire(self.owner, self.timeout)
        return self.acquired

    def __exit__(self, exc_type, exc, tb):
        if self.acquired:
            release()
            self.acquired = False
        return False


def lock(owner: str = "", timeout: float = 0.0) -> _SendLockCtx:
    """用法：with send_lock.lock("engine_b"): do_send()"""
    return _SendLockCtx(owner, timeout)


if __name__ == "__main__":
    # 冒烟测试：两个线程抢同一把锁，各自打印持锁时段
    import random

    def worker(name):
        for i in range(3):
            if try_acquire(name, 0):
                print(f"[{name}] 第{i}轮拿到锁（持有者={holder()}）")
                time.sleep(random.uniform(0.02, 0.08))
                release()
            else:
                print(f"[{name}] 第{i}轮未抢到锁，跳过")
            time.sleep(0.01)

    t1 = threading.Thread(target=worker, args=("engine_a",))
    t2 = threading.Thread(target=worker, args=("engine_b",))
    t1.start()
    t2.start()
    t1.join()
    t2.join()
    print("[OK] 发送互斥锁冒烟测试完成")