ctms/backend/app/services/permission_log_writer.py

"""权限访问日志后台写入器

使用 asyncio.Queue 缓冲日志条目，批量写入数据库。
设计原则：
- enqueue() 非阻塞，不影响请求性能
- 队列满时丢弃（监控不应拖慢被监控系统）
- 批量写入减少数据库压力
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from datetime import datetime, timezone
from typing import Any

from app.db.session import SessionLocal
from app.models.permission_access_log import PermissionAccessLog

logger = logging.getLogger("ctms.permission_log_writer")

BATCH_SIZE = 50
FLUSH_INTERVAL = 5.0
QUEUE_MAX_SIZE = 10000
WRITE_ATTEMPTS = 2
_QUEUE_STOP = object()


class PermissionLogWriter:

    def __init__(self):
        self._queue: asyncio.Queue[dict | object] = asyncio.Queue(maxsize=QUEUE_MAX_SIZE)
        self._task: asyncio.Task | None = None
        self._stopping = False
        self._started_at: datetime | None = None
        self._last_success_at: datetime | None = None
        self._last_error_at: datetime | None = None
        self._last_error_type: str | None = None
        self._accepted_entries = 0
        self._written_entries = 0
        self._dropped_entries = 0
        self._failed_batches = 0
        self._failed_entries = 0
        self._retry_count = 0

    def enqueue(self, entry: dict) -> None:
        if self._stopping:
            self._dropped_entries += 1
            logger.warning("Permission log writer is stopping, dropping entry")
            return
        try:
            self._queue.put_nowait(entry)
            self._accepted_entries += 1
        except asyncio.QueueFull:
            self._dropped_entries += 1
            logger.warning("Permission log queue full, dropping entry")

    async def start(self) -> None:
        if self._task and not self._task.done():
            return
        self._stopping = False
        self._started_at = datetime.now(timezone.utc)
        self._task = asyncio.create_task(self._flush_loop())
        logger.info("PermissionLogWriter started")

    async def stop(self) -> None:
        if self._task and not self._task.done():
            self._stopping = True
            await self._queue.put(_QUEUE_STOP)
            await self._task
        self._task = None
        logger.info("PermissionLogWriter stopped")

    async def _flush_loop(self) -> None:
        while True:
            batch, should_stop = await self._collect_batch()
            if batch:
                await self._write_batch(batch)
            if should_stop:
                break

    async def _collect_batch(self) -> tuple[list[dict], bool]:
        batch: list[dict] = []
        try:
            first = await asyncio.wait_for(self._queue.get(), timeout=FLUSH_INTERVAL)
            if first is _QUEUE_STOP:
                return batch, True
            batch.append(first)
        except asyncio.TimeoutError:
            return batch, False

        while len(batch) < BATCH_SIZE:
            try:
                item = self._queue.get_nowait()
                if item is _QUEUE_STOP:
                    return batch, True
                batch.append(item)
            except asyncio.QueueEmpty:
                break
        return batch, False

    async def _write_batch(self, batch: list[dict]) -> bool:
        for attempt in range(WRITE_ATTEMPTS):
            try:
                async with SessionLocal() as session:
                    for entry in batch:
                        log = PermissionAccessLog(
                            id=uuid.uuid4(),
                            study_id=entry["study_id"],
                            user_id=entry["user_id"],
                            endpoint_key=entry["endpoint_key"],
                            role=entry["role"],
                            allowed=entry["allowed"],
                            elapsed_ms=entry["elapsed_ms"],
                            ip_address=entry.get("ip_address"),
                            user_agent=entry.get("user_agent"),
                            client_type=entry.get("client_type"),
                            client_version=entry.get("client_version"),
                            client_platform=entry.get("client_platform"),
                            build_channel=entry.get("build_channel"),
                            build_commit=entry.get("build_commit"),
                            request_headers=entry.get("request_headers"),
                            request_snapshot=entry.get("request_snapshot"),
                            request_id=entry.get("request_id"),
                            created_at=entry.get("created_at", datetime.now(timezone.utc)),
                        )
                        session.add(log)
                    await session.commit()
            except Exception as exc:
                self._last_error_at = datetime.now(timezone.utc)
                self._last_error_type = type(exc).__name__
                if attempt + 1 < WRITE_ATTEMPTS:
                    self._retry_count += 1
                    logger.warning(
                        "Permission access log batch write failed; retrying (%d entries)",
                        len(batch),
                    )
                    await asyncio.sleep(0)
                    continue
                self._failed_batches += 1
                self._failed_entries += len(batch)
                logger.exception(
                    "Failed to write permission access log batch after retries (%d entries)",
                    len(batch),
                )
                return False
            self._written_entries += len(batch)
            self._last_success_at = datetime.now(timezone.utc)
            return True
        return False

    def stats(self) -> dict[str, Any]:
        return {
            "running": bool(self._task and not self._task.done()),
            "stopping": self._stopping,
            "queue_size": self._queue.qsize(),
            "queue_capacity": self._queue.maxsize,
            "accepted_entries": self._accepted_entries,
            "written_entries": self._written_entries,
            "dropped_entries": self._dropped_entries,
            "failed_batches": self._failed_batches,
            "failed_entries": self._failed_entries,
            "retry_count": self._retry_count,
            "started_at": self._started_at.isoformat() if self._started_at else None,
            "last_success_at": self._last_success_at.isoformat() if self._last_success_at else None,
            "last_error_at": self._last_error_at.isoformat() if self._last_error_at else None,
            "last_error_type": self._last_error_type,
        }


_writer: PermissionLogWriter | None = None


def get_log_writer() -> PermissionLogWriter | None:
    return _writer


async def start_log_writer() -> PermissionLogWriter:
    global _writer
    _writer = PermissionLogWriter()
    await _writer.start()
    return _writer


async def stop_log_writer() -> None:
    global _writer
    if _writer:
        await _writer.stop()
        _writer = None