197 lines
6.8 KiB
Python
197 lines
6.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""访问统计的内存缓冲与批量入库。
|
|
|
|
前端通过 ``POST /api/visit`` 上报页面浏览。为降低 SQLite 压力，上报先写入进程内
|
|
聚合缓冲，再由后台线程按固定周期（默认 60 秒）批量 upsert 到 ``visit`` 表。按
|
|
(日期, IP, 路由) 聚合，同一分钟的重复浏览只累加计数。
|
|
|
|
进程被强杀时最多丢失一个刷新周期的数据，属可接受范围。批量写入经 Database 的
|
|
单写线程串行执行，网络/请求路径不直接触碰数据库。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
from ..database import Database
|
|
from ..timeutil import shanghai_date
|
|
|
|
# 缓冲里每个键保存 [hits, first_at, last_at]。
|
|
_UPSERT_SQL = """
|
|
INSERT INTO visit(stat_date, ip, path, hits, first_at, last_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(stat_date, ip, path) DO UPDATE SET
|
|
hits = visit.hits + excluded.hits,
|
|
first_at = MIN(visit.first_at, excluded.first_at),
|
|
last_at = MAX(visit.last_at, excluded.last_at)
|
|
"""
|
|
|
|
|
|
class VisitCollector:
|
|
def __init__(
|
|
self,
|
|
db: Database,
|
|
*,
|
|
clock: Callable[[], float] = time.time,
|
|
flush_seconds: int = 60,
|
|
keep_days: int = 180,
|
|
max_hits_per_key: int = 0,
|
|
logger: Any = None,
|
|
) -> None:
|
|
self._db = db
|
|
self._clock = clock
|
|
self._flush_seconds = max(0, int(flush_seconds))
|
|
self._keep_days = max(1, int(keep_days))
|
|
self._max_hits = max(0, int(max_hits_per_key))
|
|
self._logger = logger
|
|
self._lock = threading.RLock()
|
|
self._buckets: dict[tuple[str, str, str], list[int]] = {}
|
|
self._last_prune_date: str | None = None
|
|
self._stop = threading.Event()
|
|
self._thread: threading.Thread | None = None
|
|
|
|
# ------------------------------------------------------------ buffering
|
|
|
|
def record(self, ip: str, path: str, *, now: float | None = None) -> bool:
|
|
"""把一次浏览累加到内存缓冲；超过单键上限时返回 ``False``。"""
|
|
if not ip or not path:
|
|
return False
|
|
ts = int(self._clock() if now is None else now)
|
|
key = (shanghai_date(ts), ip, path)
|
|
with self._lock:
|
|
bucket = self._buckets.get(key)
|
|
if bucket is None:
|
|
bucket = [0, ts, ts]
|
|
self._buckets[key] = bucket
|
|
if self._max_hits and bucket[0] >= self._max_hits:
|
|
return False
|
|
bucket[0] += 1
|
|
bucket[1] = min(bucket[1], ts)
|
|
bucket[2] = max(bucket[2], ts)
|
|
return True
|
|
|
|
def pending_rows(self) -> list[dict[str, Any]]:
|
|
"""返回未落库缓冲的快照，供 /api/stats 合并展示。"""
|
|
with self._lock:
|
|
return [
|
|
{
|
|
"stat_date": day,
|
|
"ip": ip,
|
|
"path": path,
|
|
"hits": bucket[0],
|
|
"first_at": bucket[1],
|
|
"last_at": bucket[2],
|
|
}
|
|
for (day, ip, path), bucket in self._buckets.items()
|
|
]
|
|
|
|
@property
|
|
def pending_keys(self) -> int:
|
|
with self._lock:
|
|
return len(self._buckets)
|
|
|
|
# ------------------------------------------------------------ flush
|
|
|
|
def flush(self) -> int:
|
|
"""把当前缓冲批量写入数据库，返回写入的聚合行数。"""
|
|
with self._lock:
|
|
if not self._buckets:
|
|
return 0
|
|
take = self._buckets
|
|
self._buckets = {}
|
|
rows = [
|
|
(day, ip, path, bucket[0], bucket[1], bucket[2])
|
|
for (day, ip, path), bucket in take.items()
|
|
]
|
|
|
|
def write(conn) -> None:
|
|
conn.executemany(_UPSERT_SQL, rows)
|
|
|
|
try:
|
|
self._db.transaction(write)
|
|
except Exception as exc: # noqa: BLE001 写失败不应影响服务
|
|
# 按用户容忍度可丢弃，但内存极小，这里并回缓冲等待下次重试。
|
|
with self._lock:
|
|
self._merge_back(take)
|
|
self._log("error", "visit.flush_error", error=str(exc))
|
|
return 0
|
|
self._maybe_prune()
|
|
self._log(
|
|
"info", "visit.flush",
|
|
rows=len(rows), hits=sum(row[3] for row in rows),
|
|
)
|
|
return len(rows)
|
|
|
|
def _merge_back(self, take: dict[tuple[str, str, str], list[int]]) -> None:
|
|
for key, bucket in take.items():
|
|
current = self._buckets.get(key)
|
|
if current is None:
|
|
self._buckets[key] = list(bucket)
|
|
else:
|
|
current[0] += bucket[0]
|
|
current[1] = min(current[1], bucket[1])
|
|
current[2] = max(current[2], bucket[2])
|
|
|
|
def _maybe_prune(self) -> None:
|
|
if self._keep_days <= 0:
|
|
return
|
|
today = shanghai_date(self._clock())
|
|
if today == self._last_prune_date:
|
|
return
|
|
cutoff = shanghai_date(
|
|
self._clock() - self._keep_days * 86400
|
|
)
|
|
|
|
def prune(conn) -> int:
|
|
cur = conn.execute(
|
|
"DELETE FROM visit WHERE stat_date < ?", (cutoff,)
|
|
)
|
|
return int(cur.rowcount or 0)
|
|
|
|
try:
|
|
removed = self._db.write(prune)
|
|
except Exception as exc: # noqa: BLE001
|
|
self._log("error", "visit.prune_error", error=str(exc))
|
|
return
|
|
self._last_prune_date = today
|
|
if removed:
|
|
self._log("info", "visit.prune", removed=removed, cutoff=cutoff)
|
|
|
|
# ------------------------------------------------------------ lifecycle
|
|
|
|
def start(self) -> None:
|
|
if self._flush_seconds <= 0 or self._thread is not None:
|
|
return
|
|
self._stop.clear()
|
|
self._thread = threading.Thread(
|
|
target=self._loop, name="visit-flusher", daemon=True
|
|
)
|
|
self._thread.start()
|
|
|
|
def _loop(self) -> None:
|
|
while not self._stop.wait(self._flush_seconds):
|
|
try:
|
|
self.flush()
|
|
except Exception as exc: # noqa: BLE001 后台线程自保
|
|
self._log("error", "visit.flush_loop_error", error=str(exc))
|
|
|
|
def stop(self) -> None:
|
|
"""停止后台线程并做最后一次 flush。"""
|
|
self._stop.set()
|
|
thread = self._thread
|
|
if thread is not None:
|
|
thread.join(timeout=5)
|
|
self._thread = None
|
|
try:
|
|
self.flush()
|
|
except Exception as exc: # noqa: BLE001
|
|
self._log("error", "visit.stop_flush_error", error=str(exc))
|
|
|
|
def _log(self, level: str, event: str, **fields: Any) -> None:
|
|
if self._logger is None:
|
|
return
|
|
method = getattr(self._logger, level, None)
|
|
if method is not None:
|
|
method(event, **fields) |