gengx

wechat_rpa/send_lock.py

@@ -0,0 +1,193 @@
+"""
+发送互斥锁
+==========
+双引擎并存（引擎 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] 发送互斥锁冒烟测试完成")