生产是 `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,
|
"name": channel.name,
|
||||||
"status": channel.status,
|
"status": channel.status,
|
||||||
"connected": bool(channel.bot_token_enc) and channel.status == "active",
|
"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,
|
"ilink_user_id": channel.ilink_user_id,
|
||||||
"last_seen_at": channel.last_seen_at,
|
"last_seen_at": channel.last_seen_at,
|
||||||
"last_error": channel.last_error,
|
"last_error": channel.last_error,
|
||||||
|
|||||||
@@ -11,6 +11,9 @@
|
|||||||
3. **token 失效要停**:服务端 errcode `-14` 表示登录态作废,这时停轮询、把通道标记成
|
3. **token 失效要停**:服务端 errcode `-14` 表示登录态作废,这时停轮询、把通道标记成
|
||||||
需要重新扫码,别拿死 token 一直撞墙。
|
需要重新扫码,别拿死 token 一直撞墙。
|
||||||
4. **回复失败不影响已经落库的事项**:任务记下了就是记下了,回不回得去另说。
|
4. **回复失败不影响已经落库的事项**:任务记下了就是记下了,回不回得去另说。
|
||||||
|
5. **一个通道只能有一个进程在收**:生产是 gunicorn `-w 2`,两个进程各自跑 lifespan,
|
||||||
|
不加锁就会出现两条 getupdates 同时拉同一个会话——消息被其中一个消费、回复却从另一个发出,
|
||||||
|
游标还会互相覆盖。用 Postgres 的 advisory lock 做租约:抢到的那个进程收,另一个等着接班。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -24,7 +27,7 @@ from sqlalchemy import select
|
|||||||
from sqlalchemy import text as sa_text
|
from sqlalchemy import text as sa_text
|
||||||
|
|
||||||
from ..config import get_settings
|
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 ..logging_setup import get_logger
|
||||||
from ..models import Workspace
|
from ..models import Workspace
|
||||||
from ..models_integrations import BotChannel, BotMessage
|
from ..models_integrations import BotChannel, BotMessage
|
||||||
@@ -36,6 +39,9 @@ logger = get_logger("wechat_bot")
|
|||||||
# 长轮询一次最多等这么久(服务端 75 秒超时,客户端留点余量)。
|
# 长轮询一次最多等这么久(服务端 75 秒超时,客户端留点余量)。
|
||||||
POLL_TIMEOUT_SECONDS = 80.0
|
POLL_TIMEOUT_SECONDS = 80.0
|
||||||
RETRY_DELAY_SECONDS = 5.0
|
RETRY_DELAY_SECONDS = 5.0
|
||||||
|
# 抢租约失败时隔多久再试:主进程挂了,备用进程要能在一两个呼吸之内接上。
|
||||||
|
LEASE_RETRY_SECONDS = 10.0
|
||||||
|
LEASE_NAMESPACE = "dadao.wechat_bot"
|
||||||
|
|
||||||
|
|
||||||
def site_base() -> str:
|
def site_base() -> str:
|
||||||
@@ -48,11 +54,34 @@ def site_base() -> str:
|
|||||||
return ""
|
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:
|
class BotManager:
|
||||||
"""一个进程里的轮询任务表:一个通道一个任务。"""
|
"""一个进程里的轮询任务表:一个通道一个任务。"""
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self._tasks: dict[uuid.UUID, asyncio.Task[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)
|
task = self._tasks.get(channel_id)
|
||||||
return task is not None and not task.done()
|
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:
|
async def _load(self, channel_id: uuid.UUID) -> tuple[BotChannel, str] | None:
|
||||||
@@ -112,9 +145,29 @@ class BotManager:
|
|||||||
return channel, token
|
return channel, token
|
||||||
|
|
||||||
async def _run(self, channel_id: uuid.UUID) -> None:
|
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)
|
loaded = await self._load(channel_id)
|
||||||
if loaded is None:
|
if loaded is None:
|
||||||
self._tasks.pop(channel_id, None)
|
|
||||||
return
|
return
|
||||||
channel, token = loaded
|
channel, token = loaded
|
||||||
cursor = channel.sync_buf or ""
|
cursor = channel.sync_buf or ""
|
||||||
@@ -129,7 +182,6 @@ class BotManager:
|
|||||||
except wechat_ilink.WechatError as exc:
|
except wechat_ilink.WechatError as exc:
|
||||||
if exc.stale_token:
|
if exc.stale_token:
|
||||||
await self._mark_relogin(channel_id, str(exc))
|
await self._mark_relogin(channel_id, str(exc))
|
||||||
self._tasks.pop(channel_id, None)
|
|
||||||
return
|
return
|
||||||
failures += 1
|
failures += 1
|
||||||
await self._note_error(channel_id, str(exc))
|
await self._note_error(channel_id, str(exc))
|
||||||
|
|||||||
@@ -338,3 +338,19 @@ async def test_upstream_wait_status_is_normalized(monkeypatch: pytest.MonkeyPatc
|
|||||||
monkeypatch.setattr(wechat_ilink, "_post", fake_post)
|
monkeypatch.setattr(wechat_ilink, "_post", fake_post)
|
||||||
result = await wechat_ilink.qrcode_status("qr-x")
|
result = await wechat_ilink.qrcode_status("qr-x")
|
||||||
assert result.status == "waiting"
|
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()
|
||||||
|
|||||||
Reference in New Issue
Block a user