""" 心跳监控服务 """ 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.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("重置用户心跳缓存失败")