prod/infra/monitoring/pid_manager.py

"""
Модуль для управления PID файлами процессов
Общий модуль для всех ботов в проекте
"""
import os
import sys
import signal
import atexit
import logging
from typing import Optional

logger = logging.getLogger(__name__)


class PIDManager:
    """Класс для управления PID файлами"""
    
    def __init__(self, pid_file_path: str, process_name: str = "process"):
        """
        Инициализация PID менеджера
        
        Args:
            pid_file_path: Путь к PID файлу
            process_name: Имя процесса для логирования
        """
        self.pid_file_path = pid_file_path
        self.process_name = process_name
        self.pid = os.getpid()
        
    def create_pid_file(self) -> bool:
        """
        Создание PID файла с текущим PID процесса
        
        Returns:
            bool: True если файл создан успешно, False в противном случае
        """
        try:
            # Создаем директорию если не существует
            pid_dir = os.path.dirname(self.pid_file_path)
            if pid_dir and not os.path.exists(pid_dir):
                os.makedirs(pid_dir, exist_ok=True)
            
            # Записываем PID в файл
            with open(self.pid_file_path, 'w') as f:
                f.write(str(self.pid))
            
            logger.info(f"PID файл создан для {self.process_name}: {self.pid_file_path} (PID: {self.pid})")
            
            # Регистрируем функцию очистки при завершении
            atexit.register(self.cleanup_pid_file)
            
            # Регистрируем обработчики сигналов для корректной очистки
            signal.signal(signal.SIGTERM, self._signal_handler)
            signal.signal(signal.SIGINT, self._signal_handler)
            
            return True
            
        except Exception as e:
            logger.error(f"Ошибка при создании PID файла для {self.process_name}: {e}")
            return False
    
    def cleanup_pid_file(self):
        """Удаление PID файла при завершении процесса"""
        try:
            if os.path.exists(self.pid_file_path):
                os.remove(self.pid_file_path)
                logger.info(f"PID файл удален для {self.process_name}: {self.pid_file_path}")
        except Exception as e:
            logger.error(f"Ошибка при удалении PID файла для {self.process_name}: {e}")
    
    def _signal_handler(self, signum, frame):
        """Обработчик сигналов для корректного завершения"""
        logger.info(f"Получен сигнал {signum} для {self.process_name}, очищаем PID файл...")
        self.cleanup_pid_file()
        sys.exit(0)
    
    def is_running(self) -> bool:
        """
        Проверка, запущен ли процесс с PID из файла
        
        Returns:
            bool: True если процесс запущен, False в противном случае
        """
        try:
            if not os.path.exists(self.pid_file_path):
                return False
            
            with open(self.pid_file_path, 'r') as f:
                content = f.read().strip()
                if not content:
                    return False
                
                try:
                    pid = int(content)
                    # Проверяем, существует ли процесс с таким PID
                    os.kill(pid, 0)  # Отправляем сигнал 0 для проверки существования
                    return True
                except (ValueError, OSError):
                    # PID не валидный или процесс не существует
                    return False
                    
        except Exception as e:
            logger.error(f"Ошибка при проверке PID файла для {self.process_name}: {e}")
            return False
    
    def get_pid(self) -> Optional[int]:
        """
        Получение PID из файла
        
        Returns:
            int: PID процесса или None если файл не существует или невалидный
        """
        try:
            if not os.path.exists(self.pid_file_path):
                return None
            
            with open(self.pid_file_path, 'r') as f:
                content = f.read().strip()
                if not content:
                    return None
                
                return int(content)
                
        except (ValueError, FileNotFoundError) as e:
            logger.error(f"Ошибка при чтении PID файла для {self.process_name}: {e}")
            return None


def create_pid_manager(process_name: str, project_root: str = None) -> PIDManager:
    """
    Создание PID менеджера для указанного процесса
    
    Args:
        process_name: Имя процесса (например, 'helper_bot', 'admin_bot', etc.)
        project_root: Корневая директория проекта. Если None, определяется автоматически
        
    Returns:
        PIDManager: Экземпляр PID менеджера
    """
    if project_root is None:
        # Определяем корень проекта автоматически
        current_file = os.path.abspath(__file__)
        # Поднимаемся на 2 уровня вверх от infra/monitoring/pid_manager.py
        project_root = os.path.dirname(os.path.dirname(current_file))
    
    pid_file_path = os.path.join(project_root, f"{process_name}.pid")
    
    return PIDManager(pid_file_path, process_name)


def get_bot_pid_manager(bot_name: str) -> PIDManager:
    """
    Удобная функция для создания PID менеджера для ботов
    
    Args:
        bot_name: Имя бота (например, 'helper_bot', 'admin_bot', etc.)
        
    Returns:
        PIDManager: Экземпляр PID менеджера
    """
    return create_pid_manager(bot_name)