diff --git a/backend/app/api/integrations.py b/backend/app/api/integrations.py index 728ab6d..c969227 100644 --- a/backend/app/api/integrations.py +++ b/backend/app/api/integrations.py @@ -66,7 +66,8 @@ def _channel_out(channel: BotChannel) -> dict[str, Any]: "name": channel.name, "status": channel.status, "connected": bool(channel.bot_token_enc) and channel.status == "active", - "polling": manager.running(channel.id), + # 只有抢到租约、真在拉消息的进程才算"接收中";备用进程的任务在等键,不算。 + "polling": manager.polling(channel.id), "ilink_user_id": channel.ilink_user_id, "last_seen_at": channel.last_seen_at, "last_error": channel.last_error, diff --git a/backend/app/services/wechat_bot.py b/backend/app/services/wechat_bot.py index cfa222a..5a15277 100644 --- a/backend/app/services/wechat_bot.py +++ b/backend/app/services/wechat_bot.py @@ -11,6 +11,9 @@ 3. **token 失效要停**:服务端 errcode `-14` 表示登录态作废,这时停轮询、把通道标记成 需要重新扫码,别拿死 token 一直撞墙。 4. **回复失败不影响已经落库的事项**:任务记下了就是记下了,回不回得去另说。 +5. **一个通道只能有一个进程在收**:生产是 gunicorn `-w 2`,两个进程各自跑 lifespan, + 不加锁就会出现两条 getupdates 同时拉同一个会话——消息被其中一个消费、回复却从另一个发出, + 游标还会互相覆盖。用 Postgres 的 advisory lock 做租约:抢到的那个进程收,另一个等着接班。 """ from __future__ import annotations @@ -24,7 +27,7 @@ from sqlalchemy import select from sqlalchemy import text as sa_text from ..config import get_settings -from ..db import SessionFactory, bind_tenant +from ..db import SessionFactory, bind_tenant, engine from ..logging_setup import get_logger from ..models import Workspace from ..models_integrations import BotChannel, BotMessage @@ -36,6 +39,9 @@ logger = get_logger("wechat_bot") # 长轮询一次最多等这么久(服务端 75 秒超时,客户端留点余量)。 POLL_TIMEOUT_SECONDS = 80.0 RETRY_DELAY_SECONDS = 5.0 +# 抢租约失败时隔多久再试:主进程挂了,备用进程要能在一两个呼吸之内接上。 +LEASE_RETRY_SECONDS = 10.0 +LEASE_NAMESPACE = "dadao.wechat_bot" def site_base() -> str: @@ -48,11 +54,34 @@ def site_base() -> str: return "" +async def acquire_lease(channel_id: uuid.UUID): + """抢"这个通道由我收消息"的租约;抢到返回连接(关掉就是释放),抢不到返回 None。 + + 用会话级 advisory lock:锁挂在连接上,进程崩了连接断,锁自然释放, + 不会留下"死进程占着锁没人收消息"的局面。 + """ + conn = await engine.connect() + try: + result = await conn.execute( + sa_text("SELECT pg_try_advisory_lock(hashtext(:ns), hashtext(:cid))"), + {"ns": LEASE_NAMESPACE, "cid": str(channel_id)}, + ) + if result.scalar(): + return conn + except Exception: + await conn.close() + raise + await conn.close() + return None + + class BotManager: """一个进程里的轮询任务表:一个通道一个任务。""" def __init__(self) -> None: self._tasks: dict[uuid.UUID, asyncio.Task[None]] = {} + # 抢到租约、真正在拉消息的通道;备用进程的任务在等租约,不在这里。 + self._polling: set[uuid.UUID] = set() # ── 生命周期 ────────────────────────────────────────── @@ -95,6 +124,10 @@ class BotManager: task = self._tasks.get(channel_id) return task is not None and not task.done() + def polling(self, channel_id: uuid.UUID) -> bool: + """只有真在拉消息的那个进程才算,"页面上的接收中"认这个。""" + return channel_id in self._polling + # ── 轮询主体 ────────────────────────────────────────── async def _load(self, channel_id: uuid.UUID) -> tuple[BotChannel, str] | None: @@ -112,9 +145,29 @@ class BotManager: return channel, token async def _run(self, channel_id: uuid.UUID) -> None: + lease = await self._wait_for_lease(channel_id) + if lease is None: + self._tasks.pop(channel_id, None) + return + try: + await self._poll_forever(channel_id) + finally: + self._polling.discard(channel_id) + self._tasks.pop(channel_id, None) + await lease.close() + + async def _wait_for_lease(self, channel_id: uuid.UUID): + """抢不到就等着——主进程一挂,最多 LEASE_RETRY_SECONDS 就由这里接上。""" + while True: + lease = await acquire_lease(channel_id) + if lease is not None: + self._polling.add(channel_id) + return lease + await asyncio.sleep(LEASE_RETRY_SECONDS) + + async def _poll_forever(self, channel_id: uuid.UUID) -> None: loaded = await self._load(channel_id) if loaded is None: - self._tasks.pop(channel_id, None) return channel, token = loaded cursor = channel.sync_buf or "" @@ -129,7 +182,6 @@ class BotManager: except wechat_ilink.WechatError as exc: if exc.stale_token: await self._mark_relogin(channel_id, str(exc)) - self._tasks.pop(channel_id, None) return failures += 1 await self._note_error(channel_id, str(exc)) diff --git a/backend/tests/test_wechat_bot.py b/backend/tests/test_wechat_bot.py index e70f01a..a616a99 100644 --- a/backend/tests/test_wechat_bot.py +++ b/backend/tests/test_wechat_bot.py @@ -338,3 +338,19 @@ async def test_upstream_wait_status_is_normalized(monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(wechat_ilink, "_post", fake_post) result = await wechat_ilink.qrcode_status("qr-x") assert result.status == "waiting" + + +async def test_only_one_process_holds_the_polling_lease() -> None: + """生产是 gunicorn `-w 2`:同一个通道只能有一个进程在拉消息,另一个等着接班。""" + channel_id = uuid.uuid4() + + first = await wechat_bot.acquire_lease(channel_id) + assert first is not None, "第一个进程该抢到" + + second = await wechat_bot.acquire_lease(channel_id) + assert second is None, "第二个进程不该同时抢到,否则两条轮询会互相覆盖游标" + + await first.close() # 进程退出,连接断,锁释放 + third = await wechat_bot.acquire_lease(channel_id) + assert third is not None, "前一个进程放手后,备用进程要能接上" + await third.close()