fix(integrations): 一个通道只让一个进程收消息,gunicorn 双进程不再抢同一条会话
CI / backend (push) Canceled after 0s

生产是 `gunicorn -w 2`,两个进程各自跑 lifespan,start_all() 会把同一个通道的
长轮询拉两遍:消息被其中一个消费、回复从另一个发出,游标还会互相覆盖。

改成 Postgres 会话级 advisory lock 做租约:抢到的进程收消息,另一个每 10 秒重试,
主进程一挂连接断、锁释放,备用进程立刻接上。页面上的"接收中"只认抢到租约的那个
(`manager.polling()`),不再把等锁的也算进去。

补测试 `test_only_one_process_holds_the_polling_lease`:
第一个抢到、第二个抢不到、第一个放手后第三个能接上。
This commit is contained in:
Steven
2026-11-05 23:28:35 +08:00
parent 1e8cd01b3d
commit 07b01db7cc
3 changed files with 73 additions and 4 deletions
+55 -3
View File
@@ -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))