rex/process_monitor.py
Rama master · réplica de lectura del repositorio original.
"""Process monitor for OS-level launched applications.
Tracks launched processes, detects crashes, and supports restart.
Usage::
import subprocess
from rex.process_monitor import ProcessMonitor
monitor = ProcessMonitor()
proc = subprocess.Popen(["notepad.exe"], shell=False,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
monitor.watch(proc.pid, "notepad", proc, executable="notepad.exe")
status = monitor.check(proc.pid)
if status.crashed:
restart_result = monitor.restart(proc.pid)
"""
from __future__ import annotations
import logging
import subprocess
from dataclasses import dataclass
logger = logging.getLogger(__name__)
@dataclass
class ProcessRecord:
"""Internal record of a watched process."""
pid: int
app_name: str
executable: str
process: subprocess.Popen[bytes]
auto_restart: bool = False
restart_count: int = 0
max_restarts: int = 5
@dataclass
class ProcessStatus:
"""Status snapshot for a watched process."""
pid: int
app_name: str
running: bool
crashed: bool
returncode: int | None = None
restart_count: int = 0
@dataclass
class RestartResult:
"""Result of a restart attempt."""
success: bool
app_name: str
old_pid: int
new_pid: int | None = None
error: str | None = None
class ProcessNotWatchedError(Exception):
"""Raised when a PID is not registered with the monitor."""
class RestartLimitExceededError(Exception):
"""Raised when a process has exceeded its maximum restart count."""
class ProcessMonitor:
"""Monitor launched processes for crashes and restart them on failure.
Processes must be registered with :meth:`watch` after launch. Call
:meth:`check` or :meth:`check_all` to detect crashes. Call
:meth:`restart` to re-launch a crashed process.
"""
def __init__(self) -> None:
self._records: dict[int, ProcessRecord] = {}
# ------------------------------------------------------------------
# Watch / unwatch
# ------------------------------------------------------------------
def watch(
self,
pid: int,
app_name: str,
process: subprocess.Popen[bytes],
*,
executable: str = "",
auto_restart: bool = False,
max_restarts: int = 5,
) -> None:
"""Register a process for monitoring.
Args:
pid: PID of the running process.
app_name: Logical name for the application.
process: The :class:`subprocess.Popen` handle.
executable: Executable path used to restart the process.
auto_restart: Whether to restart the process automatically on crash.
max_restarts: Maximum number of automatic restart attempts.
"""
self._records[pid] = ProcessRecord(
pid=pid,
app_name=app_name,
executable=executable,
process=process,
auto_restart=auto_restart,
max_restarts=max_restarts,
)
logger.debug("Watching process %r (pid=%d)", app_name, pid)
def unwatch(self, pid: int) -> bool:
"""Stop monitoring a process.
Args:
pid: PID to remove.
Returns:
``True`` if the process was watched and removed, ``False`` otherwise.
"""
if pid in self._records:
del self._records[pid]
logger.debug("Stopped watching pid=%d", pid)
return True
return False
def list_watched(self) -> list[int]:
"""Return PIDs of all currently watched processes."""
return list(self._records.keys())
# ------------------------------------------------------------------
# Status checks
# ------------------------------------------------------------------
def is_running(self, pid: int) -> bool:
"""Return ``True`` if the process with the given PID is still running.
Args:
pid: PID to check.
Returns:
``True`` if the process is alive (poll() returns None).
Raises:
ProcessNotWatchedError: If *pid* is not registered.
"""
record = self._records.get(pid)
if record is None:
raise ProcessNotWatchedError(
f"PID {pid} is not watched. Register it with watch() first."
)
return record.process.poll() is None
def check(self, pid: int) -> ProcessStatus:
"""Return the current status of a watched process.
A process is considered *crashed* when it has exited with a non-zero
return code (``returncode != 0`` and ``returncode is not None``).
Args:
pid: PID to check.
Returns:
:class:`ProcessStatus` snapshot.
Raises:
ProcessNotWatchedError: If *pid* is not registered.
"""
record = self._records.get(pid)
if record is None:
raise ProcessNotWatchedError(
f"PID {pid} is not watched. Register it with watch() first."
)
returncode = record.process.poll()
running = returncode is None
crashed = not running and returncode != 0
return ProcessStatus(
pid=pid,
app_name=record.app_name,
running=running,
crashed=crashed,
returncode=returncode,
restart_count=record.restart_count,
)
def check_all(self) -> list[ProcessStatus]:
"""Return status snapshots for all watched processes.
Returns:
List of :class:`ProcessStatus` objects.
"""
return [self.check(pid) for pid in list(self._records.keys())]
# ------------------------------------------------------------------
# Restart
# ------------------------------------------------------------------
def restart(self, pid: int) -> RestartResult:
"""Restart a watched (usually crashed) process.
The process is re-launched using the executable stored in the process
record. The old PID entry is replaced with the new one.
Args:
pid: PID of the process to restart.
Returns:
:class:`RestartResult`.
Raises:
ProcessNotWatchedError: If *pid* is not registered.
RestartLimitExceededError: If the process has reached *max_restarts*.
"""
record = self._records.get(pid)
if record is None:
raise ProcessNotWatchedError(
f"PID {pid} is not watched. Register it with watch() first."
)
if record.restart_count >= record.max_restarts:
raise RestartLimitExceededError(
f"Process {record.app_name!r} (pid={pid}) has exceeded "
f"its maximum restart count of {record.max_restarts}."
)
if not record.executable:
error = (
f"Cannot restart {record.app_name!r}: no executable path stored. "
"Pass executable=<path> when calling watch()."
)
logger.warning(error)
return RestartResult(
success=False,
app_name=record.app_name,
old_pid=pid,
error=error,
)
try:
new_proc = subprocess.Popen( # noqa: S603
[record.executable],
shell=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
except Exception as exc: # noqa: BLE001
error = f"Re-launch of {record.executable!r} failed: {exc}"
logger.warning("Restart of %r failed: %s", record.app_name, error)
return RestartResult(
success=False,
app_name=record.app_name,
old_pid=pid,
error=error,
)
new_pid = new_proc.pid
new_restart_count = record.restart_count + 1
logger.info(
"Restarted %r: old pid=%d -> new pid=%d (attempt %d/%d)",
record.app_name,
pid,
new_pid,
new_restart_count,
record.max_restarts,
)
# Replace old record with new one
del self._records[pid]
self._records[new_pid] = ProcessRecord(
pid=new_pid,
app_name=record.app_name,
executable=record.executable,
process=new_proc,
auto_restart=record.auto_restart,
restart_count=new_restart_count,
max_restarts=record.max_restarts,
)
return RestartResult(
success=True,
app_name=record.app_name,
old_pid=pid,
new_pid=new_pid,
)
# ---------------------------------------------------------------------------
# Module-level singleton
# ---------------------------------------------------------------------------
_monitor: ProcessMonitor | None = None
def get_process_monitor() -> ProcessMonitor:
"""Return the module-level :class:`ProcessMonitor` singleton."""
global _monitor # noqa: PLW0603
if _monitor is None:
_monitor = ProcessMonitor()
return _monitor
def set_process_monitor(monitor: ProcessMonitor | None) -> None:
"""Replace the module-level singleton (for testing)."""
global _monitor # noqa: PLW0603
_monitor = monitor
__all__ = [
"ProcessMonitor",
"ProcessNotWatchedError",
"ProcessRecord",
"ProcessStatus",
"RestartLimitExceededError",
"RestartResult",
"get_process_monitor",
"set_process_monitor",
]