216 lines
7.4 KiB
Python
216 lines
7.4 KiB
Python
"""到期与逾期提醒。
|
||
|
||
为什么是"轮询 + 数据库租约"而不是引入 Celery/APScheduler：
|
||
|
||
1. 生产是 gunicorn 多 worker，进程内定时器会各跑一遍，通知就重复了。
|
||
`pg_try_advisory_lock` 是现成的单实例租约，和微信机器人用的是同一套做法，
|
||
不额外增加运维组件。
|
||
2. 提醒的准确度要求是"几分钟内"，不是秒级；一轮扫描就是一条按
|
||
`(workspace_id, status, due_at)` 索引的查询，代价可以忽略。
|
||
3. 重复触发用 `notification.dedupe_key` 兜底：扫描器重启、时钟抖动、
|
||
人工补跑都不会把同一条提醒写两遍。
|
||
|
||
扫描本身跨租户（平台级任务），所以先在 `app.reminder_list = 'on'` 的只读口子下
|
||
把候选任务列出来，再逐个工作区绑定租户写通知——读可以跨，写必须归位。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import math
|
||
import uuid
|
||
from contextlib import suppress
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import Any
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from sqlalchemy import select, text
|
||
|
||
from ..config import get_settings
|
||
from ..db import SessionFactory, bind_tenant, engine
|
||
from ..logging_setup import get_logger
|
||
from ..models_task import Item
|
||
from .notify import notify_once
|
||
|
||
logger = get_logger("app.reminders")
|
||
|
||
_LOCK_NAMESPACE = "dadao.reminders"
|
||
# 一轮最多处理这么多条：正常量级远小于它；真到了这个量级，说明该拆任务而不是加锁。
|
||
_SCAN_LIMIT = 500
|
||
|
||
_task: asyncio.Task[None] | None = None
|
||
_stopping = asyncio.Event()
|
||
|
||
|
||
async def _candidates() -> list[dict[str, Any]]:
|
||
"""列出所有"还没完成、有截止时间、已经进入提醒窗口"的事项。"""
|
||
settings = get_settings()
|
||
now = datetime.now(timezone.utc)
|
||
horizon = now + timedelta(hours=max(1, settings.reminder_due_soon_hours))
|
||
async with SessionFactory() as session:
|
||
await session.execute(text("SELECT set_config('app.reminder_list', 'on', true)"))
|
||
rows = (
|
||
await session.execute(
|
||
select(
|
||
Item.id,
|
||
Item.workspace_id,
|
||
Item.title,
|
||
Item.status,
|
||
Item.due_at,
|
||
Item.creator_id,
|
||
Item.assignee_id,
|
||
)
|
||
.where(
|
||
Item.deleted_at.is_(None),
|
||
Item.kind != "knowledge",
|
||
Item.status.not_in(("done", "canceled")),
|
||
Item.due_at.is_not(None),
|
||
Item.due_at <= horizon,
|
||
)
|
||
.order_by(Item.due_at)
|
||
.limit(_SCAN_LIMIT)
|
||
)
|
||
).all()
|
||
return [dict(row._mapping) for row in rows]
|
||
|
||
|
||
def _local(dt: datetime) -> datetime:
|
||
return dt.astimezone(ZoneInfo(get_settings().default_timezone))
|
||
|
||
|
||
def _due_soon_message(row: dict[str, Any], now: datetime) -> tuple[str, str]:
|
||
due_at = row["due_at"]
|
||
hours = max(1, math.ceil((due_at - now).total_seconds() / 3600))
|
||
return (
|
||
f"「{row['title']}」{hours} 小时内到期",
|
||
f"截止时间 {_local(due_at):%m-%d %H:%M}，请尽快处理。",
|
||
)
|
||
|
||
|
||
def _overdue_message(row: dict[str, Any], now: datetime) -> tuple[str, str]:
|
||
due_at = row["due_at"]
|
||
late = now - due_at
|
||
days = late.days
|
||
if days >= 1:
|
||
tail = f"已经逾期 {days} 天"
|
||
else:
|
||
tail = f"已经逾期 {max(1, int(late.total_seconds() // 3600))} 小时"
|
||
return (
|
||
f"「{row['title']}」已逾期",
|
||
f"原定 {_local(due_at):%m-%d %H:%M} 完成，{tail}。",
|
||
)
|
||
|
||
|
||
async def _remind_workspace(ws_id: uuid.UUID, rows: list[dict[str, Any]], now: datetime) -> int:
|
||
sent = 0
|
||
async with SessionFactory() as session, session.begin():
|
||
await bind_tenant(session, ws_id, None)
|
||
for row in rows:
|
||
recipient = row["assignee_id"] or row["creator_id"]
|
||
if recipient is None:
|
||
continue
|
||
due_at: datetime = row["due_at"]
|
||
if due_at > now:
|
||
title, body = _due_soon_message(row, now)
|
||
kind = "task_due_soon"
|
||
dedupe = f"task_due_soon:{row['id']}:{due_at.isoformat()}"
|
||
else:
|
||
title, body = _overdue_message(row, now)
|
||
kind = "task_overdue"
|
||
today = _local(now).date().isoformat()
|
||
dedupe = f"task_overdue:{row['id']}:{due_at.isoformat()}:{today}"
|
||
if await notify_once(
|
||
session,
|
||
ws_id=ws_id,
|
||
user_id=recipient,
|
||
kind=kind,
|
||
item_id=row["id"],
|
||
title=title,
|
||
body=body,
|
||
dedupe_key=dedupe,
|
||
):
|
||
sent += 1
|
||
return sent
|
||
|
||
|
||
async def tick() -> dict[str, int]:
|
||
"""扫一轮。返回本轮的统计，供日志与测试使用。"""
|
||
settings = get_settings()
|
||
if not settings.reminders_enabled:
|
||
return {"scanned": 0, "sent": 0, "skipped": 1}
|
||
|
||
lock_conn = await engine.connect()
|
||
try:
|
||
got = (
|
||
await lock_conn.execute(
|
||
text("SELECT pg_try_advisory_lock(hashtext(:ns))"), {"ns": _LOCK_NAMESPACE}
|
||
)
|
||
).scalar()
|
||
if not got:
|
||
# 另一个 worker 正在扫，这一轮让给它。
|
||
return {"scanned": 0, "sent": 0, "skipped": 1}
|
||
|
||
now = datetime.now(timezone.utc)
|
||
rows = await _candidates()
|
||
by_ws: dict[uuid.UUID, list[dict[str, Any]]] = {}
|
||
for row in rows:
|
||
by_ws.setdefault(row["workspace_id"], []).append(row)
|
||
|
||
sent = 0
|
||
for ws_id, group in by_ws.items():
|
||
sent += await _remind_workspace(ws_id, group, now)
|
||
if sent:
|
||
logger.info(
|
||
"到期提醒已发送",
|
||
extra={"extra_fields": {"scanned": len(rows), "sent": sent}},
|
||
)
|
||
return {"scanned": len(rows), "sent": sent, "skipped": 0}
|
||
finally:
|
||
# session 级 advisory lock 挂在连接上；解锁后连接才能安全回到池里。
|
||
# 连接已断时解锁失败无需处理，连接本身也回不了池。
|
||
with suppress(Exception):
|
||
await lock_conn.execute(
|
||
text("SELECT pg_advisory_unlock(hashtext(:ns))"), {"ns": _LOCK_NAMESPACE}
|
||
)
|
||
await lock_conn.close()
|
||
|
||
|
||
async def _loop() -> None:
|
||
settings = get_settings()
|
||
while not _stopping.is_set():
|
||
try:
|
||
await tick()
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception: # pragma: no cover - 单轮失败不该把循环打死
|
||
logger.exception("到期提醒扫描失败")
|
||
try:
|
||
await asyncio.wait_for(
|
||
_stopping.wait(), timeout=max(5, settings.reminder_interval_seconds)
|
||
)
|
||
except asyncio.TimeoutError:
|
||
continue
|
||
|
||
|
||
def start() -> None:
|
||
"""启动后台提醒循环。测试环境与关闭开关时直接不启动。"""
|
||
global _task
|
||
settings = get_settings()
|
||
if settings.env == "test" or not settings.reminders_enabled:
|
||
return
|
||
if _task is not None and not _task.done():
|
||
return
|
||
_stopping.clear()
|
||
_task = asyncio.create_task(_loop(), name="dadao-reminders")
|
||
|
||
|
||
async def shutdown() -> None:
|
||
global _task
|
||
_stopping.set()
|
||
if _task is None:
|
||
return
|
||
_task.cancel()
|
||
with suppress(asyncio.CancelledError, Exception):
|
||
await _task
|
||
_task = None
|