buff/backend/app/services/heartbeat_monitor.py

"""
心跳监控服务
"""
import time
from typing import Dict, Optional
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
from app.database import get_db
from app.models.user import User
from app.models.session import Session
from app.models.system_config import SystemConfig
from app.services.session_service import SessionService
from app.core.logger import logger
from app.utils.timezone import beijing_now

class HeartbeatMonitor:
    """心跳监控服务"""

    # 内存缓存心跳状态，避免频繁数据库查询
    _heartbeat_cache: Dict[int, Dict] = {}
    _last_cleanup: float = 0

    @staticmethod
    def record_heartbeat_success(db: Session, user_id: int, session_id: str) -> None:
        """记录心跳成功"""
        try:
            cache_key = f"{user_id}:{session_id}"

            # 更新内存缓存
            if cache_key not in HeartbeatMonitor._heartbeat_cache:
                HeartbeatMonitor._heartbeat_cache[cache_key] = {
                    "user_id": user_id,
                    "session_id": session_id,
                    "failure_count": 0,
                    "last_success": time.time(),
                    "last_failure": 0,
                    "is_protected": False
                }
            else:
                cache_entry = HeartbeatMonitor._heartbeat_cache[cache_key]
                cache_entry["failure_count"] = 0  # 重置失败计数
                cache_entry["last_success"] = time.time()
                cache_entry["is_protected"] = False

            logger.debug(f"用户 {user_id} 心跳成功，失败计数已重置")

        except Exception as e:
            logger.exception("记录心跳成功失败")

    @staticmethod
    def record_heartbeat_failure(db: Session, user_id: int, session_id: str) -> bool:
        """
        记录心跳失败

        Returns:
            bool: 是否触发了保护措施
        """
        try:
            cache_key = f"{user_id}:{session_id}"

            # 使用默认的最大失败次数（可以后续从配置中获取）
            max_failures = 3

            # 更新内存缓存
            if cache_key not in HeartbeatMonitor._heartbeat_cache:
                HeartbeatMonitor._heartbeat_cache[cache_key] = {
                    "user_id": user_id,
                    "session_id": session_id,
                    "failure_count": 1,
                    "last_success": 0,
                    "last_failure": time.time(),
                    "is_protected": False
                }
            else:
                cache_entry = HeartbeatMonitor._heartbeat_cache[cache_key]
                cache_entry["failure_count"] += 1
                cache_entry["last_failure"] = time.time()

            failure_count = HeartbeatMonitor._heartbeat_cache[cache_key]["failure_count"]

            logger.warning(f"用户 {user_id} 心跳失败 {failure_count}/{max_failures}")

            # 检查是否需要触发保护措施
            if failure_count >= max_failures:
                HeartbeatMonitor._execute_protection_measures(db, user_id, session_id)
                HeartbeatMonitor._heartbeat_cache[cache_key]["is_protected"] = True
                return True

            return False

        except Exception as e:
            logger.exception("记录心跳失败异常")
            return False

    @staticmethod
    def _execute_protection_measures(db: Session, user_id: int, session_id: str) -> None:
        """执行保护措施"""
        try:
            logger.warning(f"执行用户 {user_id} 的保护措施")

            # 1. 停用所有相关会话
            deactivated_count = SessionService.deactivate_user_sessions(db, user_id)
            if deactivated_count > 0:
                logger.info(f"已停用用户 {user_id} 的 {deactivated_count} 个会话")

            # 2. 可以在这里添加其他保护措施
            # 例如：通知其他相关服务、记录安全事件等

            # 3. 记录保护措施执行
            logger.warning(f"用户 {user_id} 因心跳失败过多，已执行保护措施")

        except Exception as e:
            logger.exception("执行保护措施失败")

    @staticmethod
    def get_heartbeat_status(user_id: int, session_id: str) -> Dict:
        """获取心跳状态"""
        cache_key = f"{user_id}:{session_id}"
        return HeartbeatMonitor._heartbeat_cache.get(cache_key, {
            "user_id": user_id,
            "session_id": session_id,
            "failure_count": 0,
            "last_success": 0,
            "last_failure": 0,
            "is_protected": False
        })

    @staticmethod
    def cleanup_cache() -> int:
        """
        清理过期的缓存数据

        Returns:
            int: 清理的条目数量
        """
        try:
            current_time = time.time()
            expired_keys = []

            # 查找过期条目（超过1小时没有活动的）
            for cache_key, cache_entry in HeartbeatMonitor._heartbeat_cache.items():
                last_activity = max(cache_entry["last_success"], cache_entry["last_failure"])
                if current_time - last_activity > 3600:  # 1小时
                    expired_keys.append(cache_key)

            # 删除过期条目
            for key in expired_keys:
                del HeartbeatMonitor._heartbeat_cache[key]

            if expired_keys:
                logger.info(f"清理了 {len(expired_keys)} 个过期的缓存条目")

            return len(expired_keys)

        except Exception as e:
            logger.exception("清理缓存失败")
            return 0

    @staticmethod
    def get_monitoring_stats() -> Dict:
        """获取监控统计信息"""
        try:
            total_users = len(HeartbeatMonitor._heartbeat_cache)
            protected_users = sum(1 for entry in HeartbeatMonitor._heartbeat_cache.values()
                                if entry["is_protected"])
            failing_users = sum(1 for entry in HeartbeatMonitor._heartbeat_cache.values()
                               if entry["failure_count"] > 0)

            return {
                "total_monitored_users": total_users,
                "protected_users": protected_users,
                "failing_users": failing_users,
                "healthy_users": total_users - failing_users,
                "cache_size": len(HeartbeatMonitor._heartbeat_cache)
            }

        except Exception as e:
            logger.exception("获取监控统计失败")
            return {}

    @staticmethod
    def reset_user_heartbeat(user_id: int, session_id: str) -> None:
        """重置用户的缓存状态（用于重新连接时）"""
        try:
            cache_key = f"{user_id}:{session_id}"
            if cache_key in HeartbeatMonitor._heartbeat_cache:
                del HeartbeatMonitor._heartbeat_cache[cache_key]
                logger.info(f"已重置用户 {user_id} 的心跳缓存")

        except Exception as e:
            logger.exception("重置用户心跳缓存失败")