benchlab-pytools 3.0.0__py3-none-any.whl

benchlab/core/process_manager.py

@@ -0,0 +1,356 @@
+"""
+Process Manager for BENCHLAB
+
+Structured subprocess management for the infrastructure services
+(FastAPI server, MQTT publisher, embedded amqtt broker).
+
+Replaces the ad-hoc ``subprocess.Popen`` + env-var PID pattern in main.py
+with proper lifecycle management: start, health-check, stop, restart,
+and graceful shutdown of all managed processes.
+
+Usage:
+    pm = ProcessManager.get_instance()
+
+    # Start a managed service
+    pm.start_service(
+        name="fastapi",
+        cmd=["python", "-m", "benchlab", "fastapi"],
+        health_check=lambda: _ping_http("http://127.0.0.1:8000/health"),
+        timeout=20,
+    )
+
+    # Stop a specific service
+    pm.stop_service("fastapi")
+
+    # Stop everything (called on tool exit)
+    pm.shutdown_all()
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import subprocess
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional
+
+logger = logging.getLogger("benchlab.core.process_manager")
+
+
+# ---------------------------------------------------------------------------
+# Data classes
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ManagedProcess:
+    """State for a single managed subprocess."""
+
+    name: str
+    cmd: list[str]
+    process: Optional[subprocess.Popen] = None
+    health_check: Optional[Callable[[], bool]] = None
+    stdout_log: str = ""
+    stderr_log: str = ""
+    started_at: float = 0.0
+
+
+# ---------------------------------------------------------------------------
+# Singleton
+# ---------------------------------------------------------------------------
+
+
+class ProcessManager:
+    """Thread-safe singleton that manages infrastructure service lifecycles."""
+
+    _instance: Optional["ProcessManager"] = None
+    _lock = threading.Lock()
+
+    def __init__(self) -> None:
+        self._services: Dict[str, ManagedProcess] = {}
+        self._mutex = threading.Lock()
+
+    # -- singleton helpers ---------------------------------------------------
+
+    @classmethod
+    def get_instance(cls) -> "ProcessManager":
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset singleton (primarily for testing)."""
+        with cls._lock:
+            inst = cls._instance
+            if inst is not None:
+                inst.shutdown_all()
+                cls._instance = None
+
+    # -- service management --------------------------------------------------
+
+    def start_service(
+        self,
+        name: str,
+        cmd: list[str] | str,
+        health_check: Optional[Callable[[], bool]] = None,
+        timeout: float = 20.0,
+        cwd: Optional[str] = None,
+        env: Optional[dict] = None,
+    ) -> bool:
+        """Start a service and wait for it to become healthy.
+
+        Args:
+            name: Human-readable service label (must be unique).
+            cmd: Command string or list of arguments.
+            health_check: Callable returning True when the service is ready.
+                          If ``None``, success is determined by the process
+                          staying alive for *timeout* seconds.
+            timeout: Seconds to wait for health_check to succeed.
+            cwd: Working directory for the subprocess.
+            env: Environment variables (defaults to os.environ copy).
+
+        Returns:
+            ``True`` if the service started and passed the health check.
+        """
+        with self._mutex:
+            if name in self._services:
+                existing = self._services[name]
+                logger.warning(
+                    "Service %s already registered — stopping it first", name)
+                self._stop_process(existing)
+
+        if isinstance(cmd, str):
+            cmd_list = [cmd]
+        else:
+            cmd_list = cmd
+
+        mp = ManagedProcess(name=name, cmd=cmd_list, health_check=health_check)
+
+        # Use shell execution when cmd is a single string (e.g. inline Python
+        # one-liner)
+        use_shell = isinstance(cmd, str)
+
+        log_file_prefix = os.path.join(
+            os.path.dirname(
+                os.path.abspath(__file__)),
+            "..",
+            "logs",
+            f"svc_{name}")
+        os.makedirs(os.path.dirname(log_file_prefix), exist_ok=True)
+
+        stdout_path = f"{log_file_prefix}_stdout.log"
+        stderr_path = f"{log_file_prefix}_stderr.log"
+
+        try:
+            env = env or os.environ.copy()
+            # Open log files and store file handles for later cleanup
+            stdout_file = open(stdout_path, "w", encoding="utf-8")
+            stderr_file = open(stderr_path, "w", encoding="utf-8")
+            mp.process = subprocess.Popen(
+                cmd_list,
+                stdout=stdout_file,
+                stderr=stderr_file,
+                cwd=cwd,
+                env=env,
+                shell=use_shell,
+                # On Linux, ensure the process doesn't inherit unnecessary file
+                # descriptors
+                close_fds=True,
+            )
+            # Store file handles for cleanup when service stops
+            mp._stdout_file = stdout_file
+            mp._stderr_file = stderr_file
+        except Exception as exc:
+            logger.error("Failed to start service %s: %s", name, exc)
+            return False
+
+        mp.started_at = time.time()
+        logger.info("Started service %s (PID %d)", name, mp.process.pid)
+
+        # Wait for health check
+        healthy = self._wait_for_health(mp, timeout)
+        if not healthy:
+            # Capture log content for diagnostics
+            try:
+                with open(stderr_path, "r", encoding="utf-8") as f:
+                    mp.stderr_log = f.read()
+                with open(stdout_path, "r", encoding="utf-8") as f:
+                    mp.stdout_log = f.read()
+            except OSError:
+                pass
+            logger.warning(
+                "Service %s failed health check within %.0fs",
+                name,
+                timeout)
+            with self._mutex:
+                self._services[name] = mp
+            return False
+
+        with self._mutex:
+            self._services[name] = mp
+        logger.info("Service %s is healthy", name)
+        return True
+
+    def stop_service(self, name: str) -> bool:
+        """Stop a specific service (graceful SIGTERM → wait → SIGKILL)."""
+        with self._mutex:
+            mp = self._services.pop(name, None)
+        if mp is None:
+            logger.debug("Service %s not found", name)
+            return False
+        return self._stop_process(mp)
+
+    def get_service(self, name: str) -> Optional[ManagedProcess]:
+        with self._mutex:
+            return self._services.get(name)
+
+    def list_services(self) -> list[str]:
+        with self._mutex:
+            return list(self._services.keys())
+
+    def is_running(self, name: str) -> bool:
+        mp = self.get_service(name)
+        if mp is None or mp.process is None:
+            return False
+        return mp.process.poll() is None
+
+    def shutdown_all(self) -> None:
+        """Stop all managed services.  Safe to call multiple times."""
+        with self._mutex:
+            names = list(self._services.keys())
+        for name in names:
+            self.stop_service(name)
+
+    # -- internals -----------------------------------------------------------
+
+    def _wait_for_health(self, mp: ManagedProcess, timeout: float) -> bool:
+        """Block until health_check passes or timeout.  Also fails if the
+        process dies early."""
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            # Check if process crashed
+            ret = mp.process.poll()
+            if ret is not None:
+                logger.error(
+                    "Service %s exited unexpectedly with code %d",
+                    mp.name,
+                    ret)
+                try:
+                    # try to capture the log
+                    log_base = os.path.join(
+                        os.path.dirname(
+                            os.path.abspath(__file__)),
+                        "..",
+                        "logs",
+                        f"svc_{
+                            mp.name}")
+                    with open(f"{log_base}_stderr.log", "r",
+                              encoding="utf-8") as f:
+                        mp.stderr_log = f.read()
+                except OSError:
+                    pass
+                return False
+
+            if mp.health_check is None:
+                # No health check → consider healthy after 1 second alive
+                if time.time() - mp.started_at >= 1.0:
+                    return True
+            elif mp.health_check():
+                return True
+
+            time.sleep(0.5)
+
+        logger.error(
+            "Service %s health check timed out after %.0fs",
+            mp.name,
+            timeout)
+        return False
+
+    def _stop_process(
+            self,
+            mp: ManagedProcess,
+            graceful_timeout: float = 5.0) -> bool:
+        """Terminate or kill a single ManagedProcess."""
+        proc = mp.process
+        if proc is None:
+            # Clean up any open file handles
+            self._close_log_files(mp)
+            return True
+
+        if proc.poll() is not None:
+            logger.info(
+                "Service %s already exited (code %d)",
+                mp.name,
+                proc.returncode)
+            self._close_log_files(mp)
+            return True
+
+        # Graceful: SIGTERM / taskkill
+        logger.info(
+            "Stopping service %s (PID %d) gracefully.",
+            mp.name,
+            proc.pid)
+        self._send_terminate(proc)
+
+        try:
+            proc.wait(timeout=graceful_timeout)
+            logger.info("Service %s stopped gracefully", mp.name)
+            self._close_log_files(mp)
+            return True
+        except subprocess.TimeoutExpired:
+            pass
+
+        # Forceful: SIGKILL / taskkill /F
+        logger.warning("Service %s did not stop gracefully — killing", mp.name)
+        try:
+            self._send_kill(proc)
+            proc.wait(timeout=3)
+            logger.info("Service %s killed", mp.name)
+            self._close_log_files(mp)
+            return True
+        except Exception:
+            logger.error("Failed to kill service %s", mp.name)
+            self._close_log_files(mp)
+            return False
+
+    @staticmethod
+    def _close_log_files(mp: ManagedProcess) -> None:
+        """Close any open log file handles to prevent file descriptor leaks."""
+        for attr in ('_stdout_file', '_stderr_file'):
+            f = getattr(mp, attr, None)
+            if f is not None:
+                try:
+                    f.flush()
+                    f.close()
+                except Exception:
+                    pass
+                finally:
+                    setattr(mp, attr, None)
+
+    @staticmethod
+    def _send_terminate(proc: subprocess.Popen) -> None:
+        if sys.platform == "win32":
+            subprocess.run(
+                ["taskkill", "/T", "/F", "/PID", str(proc.pid)],
+                capture_output=True,
+            )
+        else:
+            import signal
+            os.kill(proc.pid, signal.SIGTERM)
+
+    @staticmethod
+    def _send_kill(proc: subprocess.Popen) -> None:
+        if sys.platform == "win32":
+            subprocess.run(
+                ["taskkill", "/T", "/F", "/PID", str(proc.pid)],
+                capture_output=True,
+            )
+        else:
+            import signal
+            os.kill(proc.pid, signal.SIGKILL)

benchlab/core/retry.py

@@ -0,0 +1,85 @@
+"""Simple retry utilities used across Benchlab data sources.
+
+The library originally performed a single attempt to connect to a serial
+port, FastAPI server or MQTT broker.  In production environments a
+transient failure (e.g. the device is temporarily busy or the network
+is momentarily unavailable) should not cause the whole tool to abort.
+
+This module provides a lightweight ``RetryPolicy`` dataclass and a
+``retry`` decorator that can be applied to any callable.  The decorator
+retries the wrapped function according to the policy and logs each
+attempt using the module‑level ``logger``.
+"""
+
+from __future__ import annotations
+
+import time
+import logging
+from dataclasses import dataclass
+from typing import Callable, TypeVar, Any, cast
+
+logger = logging.getLogger("benchlab.core.retry")
+
+T = TypeVar("T")
+
+
+@dataclass
+class RetryPolicy:
+    """Configuration for retry behaviour.
+
+    Attributes
+    ----------
+    max_retries: int
+        Maximum number of attempts (the initial call counts as the first
+        attempt). ``0`` means no retries.
+    backoff_factor: float
+        Multiplier applied to the base delay after each failure. The
+        actual sleep time is ``base_delay * (backoff_factor ** attempt)``.
+    base_delay: float
+        Initial delay in seconds before the first retry.
+    allowed_exceptions: tuple[type[BaseException], ...]
+        Exceptions that trigger a retry. All other exceptions are raised
+        immediately.
+    """
+
+    max_retries: int = 3
+    backoff_factor: float = 2.0
+    base_delay: float = 0.5
+    allowed_exceptions: tuple[type[BaseException], ...] = (Exception,)
+
+
+def retry(
+        policy: RetryPolicy) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """Decorator that retries a function according to *policy*.
+
+    The wrapped function is called up to ``policy.max_retries + 1``
+    times.  If the function succeeds, its result is returned immediately.
+    If it raises an exception listed in ``policy.allowed_exceptions`` the
+    decorator sleeps for an exponentially‑increasing delay and retries.
+    After exhausting all attempts the last exception is re‑raised.
+    """
+
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        def wrapper(*args: Any, **kwargs: Any) -> T:
+            attempt = 0
+            while True:
+                try:
+                    return cast(T, func(*args, **kwargs))
+                except policy.allowed_exceptions as exc:
+                    if attempt >= policy.max_retries:
+                        logger.error(
+                            f"{func.__name__} failed after "
+                            f"{attempt + 1} attempts: {exc}"
+                        )
+                        raise
+                    delay = policy.base_delay * \
+                        (policy.backoff_factor ** attempt)
+                    logger.warning(
+                        f"{func.__name__} attempt {attempt + 1} failed "
+                        f"({exc}); retrying in {delay:.2f}s"
+                    )
+                    time.sleep(delay)
+                    attempt += 1
+        return wrapper
+
+    return decorator

benchlab/core/shared_serial.py

@@ -0,0 +1,235 @@
+"""
+Shared Serial Connection Manager for BENCHLAB
+
+Provides a singleton-based serial connection pool that prevents
+multiple services (FastAPI, MQTT, etc.) from competing for the
+same physical serial port. (BUG-1.4)
+"""
+
+import logging
+import threading
+import time
+from typing import Any, Dict, Optional
+
+import serial
+
+logger = logging.getLogger("benchlab.core.shared_serial")
+
+# BENCHLAB devices communicate over USB CDC at a fixed 115200 baud.
+# benchlab_pycore.core.serial_io never exposed a connection-opening helper
+# (its own docs just call serial.Serial(port, 115200, timeout=1) directly) —
+# this wraps that same pattern so callers get a consistent
+# "None on failure" contract instead of talking to pyserial directly.
+_BENCHLAB_BAUDRATE = 115200
+_BENCHLAB_TIMEOUT = 1
+
+
+def open_serial_connection(port: str) -> Optional[serial.Serial]:
+    """Open a serial connection to a BENCHLAB device.
+
+    Returns None if the port can't be opened.
+    """
+    try:
+        return serial.Serial(
+            port,
+            _BENCHLAB_BAUDRATE,
+            timeout=_BENCHLAB_TIMEOUT)
+    except (serial.SerialException, OSError) as e:
+        logger.warning("Failed to open serial port %s: %s", port, e)
+        return None
+
+
+class SharedSerialManager:
+    """Thread-safe singleton that manages serial port ownership.
+
+    Use this to ensure only one reader owns a serial port at a time,
+    while allowing multiple consumers to receive telemetry data.
+
+    Usage:
+        mgr = SharedSerialManager.get_instance()
+        conn = mgr.acquire_connection("COM3")
+        if conn:
+            data = read_sensors(conn.ser)
+        mgr.release_connection("COM3")
+    """
+
+    _instance = None
+    _instance_lock = threading.Lock()
+
+    def __init__(self):
+        # port -> {ser, ref_count, lock, owner}
+        self._connections: Dict[str, Dict[str, Any]] = {}
+        self._lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls) -> "SharedSerialManager":
+        """Get or create the singleton instance."""
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        """Reset singleton (primarily for testing)."""
+        with cls._instance_lock:
+            if cls._instance is not None:
+                cls._instance._shutdown()
+                cls._instance = None
+
+    def acquire_connection(
+        self,
+        port: str,
+        open_func=None,
+        timeout: float = 5.0
+    ) -> Optional["SharedConnection"]:
+        """Acquire a shared serial connection to a port.
+
+        Args:
+            port: Serial port identifier (e.g., "COM3", "/dev/ttyUSB0")
+            open_func: Callback to open the serial port if not already open.
+                       Signature: open_func(port) -> serial_connection
+            timeout: How long to wait if port is being opened by another thread
+
+        Returns:
+            SharedConnection wrapper, or None if port can't be opened
+        """
+        if open_func is None:
+            open_func = open_serial_connection
+
+        start = time.time()
+        while time.time() - start < timeout:
+            with self._lock:
+                # Check if already open
+                if port in self._connections:
+                    entry = self._connections[port]
+                    entry["ref_count"] += 1
+                    logger.debug(
+                        "Reusing shared connection to %s (ref_count=%d)",
+                        port,
+                        entry["ref_count"])
+                    return SharedConnection(self, port, entry["ser"])
+
+                # Try to claim it
+                self._connections[port] = {
+                    "ser": None,
+                    "ref_count": 1,
+                    "lock": threading.Lock(),
+                    "owner": threading.current_thread().name,
+                }
+
+            # Open outside the main lock to avoid blocking
+            try:
+                ser = open_func(port)
+                if ser:
+                    with self._lock:
+                        self._connections[port]["ser"] = ser
+
+                    logger.info("Opened new shared connection to %s", port)
+                    return SharedConnection(self, port, ser)
+                else:
+                    with self._lock:
+                        del self._connections[port]
+                    logger.warning("Failed to open serial port %s", port)
+                    return None
+            except Exception as e:
+                with self._lock:
+                    del self._connections[port]
+                logger.warning("Exception opening port %s: %s", port, e)
+                time.sleep(0.5)
+                continue
+
+        logger.error("Timeout acquiring connection to %s", port)
+        return None
+
+    def release_connection(self, port: str) -> None:
+        """Release a shared serial connection.
+
+        When ref_count reaches 0, the connection is closed.
+        """
+        with self._lock:
+            if port not in self._connections:
+                logger.warning(
+                    "Attempted to release non-existent connection on %s", port)
+                return
+
+            entry = self._connections[port]
+            entry["ref_count"] -= 1
+
+            if entry["ref_count"] <= 0:
+                # Last user released - close the port
+                ser = entry["ser"]
+                del self._connections[port]
+
+                if ser:
+                    try:
+                        ser.close()
+                        logger.info("Closed shared connection to %s", port)
+                    except Exception:
+                        pass
+
+    def is_port_owned(self, port: str) -> bool:
+        """Check if a port currently has an open shared connection."""
+        with self._lock:
+            return port in self._connections
+
+    def get_active_ports(self) -> list:
+        """Get list of ports with active shared connections."""
+        with self._lock:
+            return list(self._connections.keys())
+
+    def _shutdown(self) -> None:
+        """Close all connections (called by reset() or cleanup)."""
+        with self._lock:
+            for port, entry in list(self._connections.items()):
+                ser = entry["ser"]
+                if ser:
+                    try:
+                        ser.close()
+                    except Exception:
+                        pass
+            self._connections.clear()
+
+
+class SharedConnection:
+    """Wrapper around a shared serial connection.
+
+    Use `conn.ser` to access the underlying serial connection.
+    Call `conn.release()` when done.
+    """
+
+    def __init__(self, manager: SharedSerialManager, port: str, ser):
+        self._manager = manager
+        self._port = port
+        self._released = False
+        self.ser = ser  # Direct access to the serial connection
+
+    @property
+    def port(self) -> str:
+        return self._port
+
+    @property
+    def is_released(self) -> bool:
+        return self._released
+
+    def release(self) -> None:
+        """Release this reference to the shared connection."""
+        if not self._released:
+            self._released = True
+            self._manager.release_connection(self._port)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.release()
+        return False
+
+    def __del__(self):
+        # Safety net - don't leak connections
+        if not self._released:
+            logger.warning(
+                "SharedConnection %s was garbage-collected without release()",
+                self._port)
+            self.release()