生产是 `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:
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user