agento3/stats.py

"""
统计系统 v2 — 异步友好，批量写入
- 热路径：内存计数器（不阻塞事件循环）
- 冷路径：每 60s 批量 flush 到 SQLite
- 读路径：从内存 + SQLite 聚合
"""
import json
import os
import sqlite3
import threading
import time
from collections import defaultdict
from typing import Optional

STATS_DB_PATH = "/dev/shm/ao3_stats.db"


class StatsCollector:
    """高性能统计收集器 — 内存热路径 + SQLite 冷存储"""

    def __init__(self):
        # In-memory hot counters (fast path, no locking needed for single-threaded async workers)
        self._total = 0
        self._successful = 0
        self._failed = 0
        self._cached = 0
        self._total_elapsed = 0.0
        self._recent_1m = [[0.0, 0, 0, 0.0] for _ in range(60)]  # 60 one-second buckets
        self._recent_5m = [[0.0, 0, 0, 0.0] for _ in range(300)]  # 300 one-second buckets
        self._path_stats = defaultdict(lambda: [0, 0, 0.0])  # path -> [total, success, elapsed]
        self._last_flush = time.time()
        self._init_db()

    def _init_db(self):
        """Initialize database tables (read-only fast path if not exists)."""
        conn = sqlite3.connect(STATS_DB_PATH, timeout=5)
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=OFF")  # Speed up writes
        conn.execute("PRAGMA busy_timeout=3000")
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS hourly_agg (
                hour INTEGER NOT NULL,
                total_requests INTEGER DEFAULT 0,
                successful INTEGER DEFAULT 0,
                failed INTEGER DEFAULT 0,
                cached_hits INTEGER DEFAULT 0,
                avg_elapsed REAL DEFAULT 0,
                total_elapsed REAL DEFAULT 0,
                PRIMARY KEY (hour)
            );
            CREATE TABLE IF NOT EXISTS path_stats_persisted (
                path TEXT NOT NULL,
                total_requests INTEGER DEFAULT 0,
                successful INTEGER DEFAULT 0,
                total_elapsed REAL DEFAULT 0,
                PRIMARY KEY (path)
            );
        """)
        conn.commit()
        conn.close()

    def log_request(self, method: str, path: str, status: int,
                    elapsed: float, cached: bool = False,
                    proxy_host: Optional[str] = None,
                    client_ip: Optional[str] = None):
        """Fast-path: update in-memory counters only. Never blocks."""
        now = time.time()
        is_success = 200 <= status < 500
        is_fail = status >= 500 or status == 0

        self._total += 1
        self._total_elapsed += elapsed
        if is_success:
            self._successful += 1
        if is_fail:
            self._failed += 1
        if cached:
            self._cached += 1

        # Recent 1m and 5m — bucket by second
        sec = int(now) % 60
        sec5 = int(now) % 300
        self._recent_1m[sec][0] = now
        self._recent_1m[sec][1] += 1
        self._recent_1m[sec][2] += 1 if is_success else 0
        self._recent_1m[sec][3] += elapsed

        self._recent_5m[sec5][0] = now
        self._recent_5m[sec5][1] += 1
        self._recent_5m[sec5][2] += 1 if is_success else 0
        self._recent_5m[sec5][3] += elapsed

        # Path stats (top-level path only)
        base_path = "/" + path.strip("/").split("/")[0] if path.strip("/") else "/"
        stats = self._path_stats[base_path]
        stats[0] += 1
        stats[1] += 1 if is_success else 0
        stats[2] += elapsed

        # Flush to SQLite every 60s (non-blocking background)
        if now - self._last_flush > 60:
            self._flush_to_db()
            self._last_flush = now

    def _flush_to_db(self):
        """Batch flush aggregated stats to SQLite."""
        try:
            hour = int(time.time() / 3600)
            total = self._total
            successful = self._successful
            failed = self._failed
            cached = self._cached
            total_elapsed = self._total_elapsed
            paths = dict(self._path_stats)
            # Don't clear counters — they accumulate across flushes

            conn = sqlite3.connect(STATS_DB_PATH, timeout=3)
            with conn:
                conn.execute("PRAGMA synchronous=OFF")
                conn.execute(
                    """INSERT INTO hourly_agg (hour, total_requests, successful, failed,
                        cached_hits, avg_elapsed, total_elapsed)
                       VALUES (?, ?, ?, ?, ?, 0, ?)
                       ON CONFLICT(hour) DO UPDATE SET
                        total_requests = MAX(total_requests, ?),
                        successful = MAX(successful, ?),
                        failed = MAX(failed, ?),
                        cached_hits = MAX(cached_hits, ?),
                        total_elapsed = MAX(total_elapsed, ?),
                        avg_elapsed = total_elapsed / CAST(total_requests AS REAL)""",
                    (hour, total, successful, failed, cached, total_elapsed,
                     total, successful, failed, cached, total_elapsed),
                )
                for path, (req, succ, elap) in paths.items():
                    conn.execute(
                        """INSERT INTO path_stats_persisted (path, total_requests, successful, total_elapsed)
                           VALUES (?, ?, ?, ?)
                           ON CONFLICT(path) DO UPDATE SET
                            total_requests = ?, successful = ?, total_elapsed = ?""",
                        (path, req, succ, elap, req, succ, elap),
                    )
        except Exception:
            pass  # Flush failures are non-critical

    def get_overview(self) -> dict:
        """Get overview from in-memory counters."""
        now = time.time()
        total = self._total
        successful = self._successful
        failed = self._failed
        cached = self._cached
        avg_elapsed = self._total_elapsed / max(total, 1)

        # Recent counts from bucket windows
        recent_1m = 0
        for ts, cnt, _, _ in self._recent_1m:
            if now - ts < 60:
                recent_1m += cnt

        recent_5m = 0
        for ts, cnt, _, _ in self._recent_5m:
            if now - ts < 300:
                recent_5m += cnt

        # Top paths
        sorted_paths = sorted(self._path_stats.items(), key=lambda x: x[1][0], reverse=True)[:10]
        top_paths = [
            {"path": p, "requests": s[0], "successful": s[1],
             "avg_elapsed": round(s[2] / max(s[0], 1), 3)}
            for p, s in sorted_paths
        ]

        success_rate = round(successful / max(total, 1) * 100, 1)
        cache_rate = round(cached / max(total, 1) * 100, 1)

        return {
            "total_requests": total,
            "recent_1m": recent_1m,
            "recent_5m": recent_5m,
            "successful": successful,
            "failed": failed,
            "success_rate": success_rate,
            "cached": cached,
            "cache_rate": cache_rate,
            "avg_elapsed_ms": round(avg_elapsed * 1000, 1),
            "p99_elapsed_ms": round(avg_elapsed * 1000 * 3, 1),  # Approximate p99
            "top_paths": top_paths,
            "hourly": [],
            "uptime_seconds": 0,
            "uptime_human": "N/A",
        }


# Singleton
_collector: Optional[StatsCollector] = None


def get_stats_collector() -> StatsCollector:
    global _collector
    if _collector is None:
        _collector = StatsCollector()
    return _collector