maqet 0.0.1.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

maqet/handlers/stage.py

@@ -0,0 +1,196 @@
+# from subprocess import check_output
+import os
+import pty
+import signal
+import subprocess
+from time import sleep
+
+from maqet.handlers.base import Handler, HandlerError
+from maqet.logger import LOG
+from maqet.qmp import commands as qmp_commands
+from maqet.qmp.keyboard import KeyboardEmulator as kb
+
+
+class StageHandler(Handler):
+    """
+    Handles execution of pipeline stages. Methods are tasks
+    """
+
+
+@StageHandler.method
+def launch(state: dict):
+    state.vm.launch()
+
+
+@StageHandler.method
+def shutdown(state,
+             hard: bool = False,
+             timeout: int = 30) -> None:
+    state.vm.shutdown(hard=hard, timeout=timeout)
+
+
+@StageHandler.method
+def wait_for_input(state, prompt: str = ""):
+    input(prompt)
+
+
+@StageHandler.method
+def wait_for_shutdown(state):
+    while state.vm.is_running():
+        sleep(3)
+
+
+@StageHandler.method
+def qmp_key(state, keys: [str],
+            hold_time: int = 1,
+            **kwargs):
+    command = kb.press_keys(*keys, hold_time=hold_time)
+    qmp_run_command(state.vm, **command)
+
+
+def qmp_run_command(vm,
+                    command: str,
+                    arguments: dict = None):
+    r = vm.qmp(cmd=command,
+               args_dict=arguments)
+    LOG.info(f"QMP: {command} {arguments} {r}")
+
+
+@StageHandler.method
+def qmp_type(state, text, type_delay: int = 10,
+             hold_time: int = 1, ** kwargs):
+    for command in kb.type_string(string=text,
+                                  hold_time=hold_time):
+        qmp_run_command(state.vm, **command)
+        sleep(type_delay/1000)
+
+
+@StageHandler.method
+def qmp(state,
+        cmd: str = None,
+        args_dict: dict = None,
+        command: str = None,
+        arguments: dict = None):
+    # Support both old and new parameter names for backward compatibility
+    command = command or cmd
+    arguments = arguments or args_dict
+    qmp_run_command(
+        vm=state.vm,
+        command=command,
+        arguments=arguments
+    )
+
+
+@StageHandler.method
+def qmp_screendump(state, filename: str):
+    command = qmp_commands.qmp_screendump(filename)
+    qmp_run_command(state.vm, **command)
+
+
+@StageHandler.method
+def qmp_stop(state):
+    command = qmp_commands.qmp_stop()
+    qmp_run_command(state.vm, **command)
+
+
+@StageHandler.method
+def qmp_cont(state):
+    command = qmp_commands.qmp_cont()
+    qmp_run_command(state.vm, **command)
+
+
+@StageHandler.method
+def qmp_pmemsave(state, filename: str):
+    # Get the memory size from the machine's memory attribute
+    size_bytes = state.vm.memory
+    # Set address to 0 to dump all memory
+    address = 0
+    command = qmp_commands.qmp_pmemsave(address, size_bytes, filename)
+    qmp_run_command(state.vm, **command)
+
+
+@StageHandler.method
+def wait(state, time: float):
+    sleep(float(time))
+
+
+@StageHandler.method
+def bash(state, script: str, silent=False, blocking=True, fatal=True, **kwargs):
+    LOG.debug(f"Executing bash script: {script}")
+
+    if not blocking:
+        raise HandlerError(
+            "Non-blocking bash tasks are not currently supported.")
+
+    process = None
+    try:
+        # For silent execution, redirect stdout and stderr
+        stdout = subprocess.DEVNULL if silent else None
+        stderr = subprocess.DEVNULL if silent else None
+
+        # Use Popen for better control over signal handling
+        # Don't create a new process group - let signals propagate naturally
+        process = subprocess.Popen(
+            script,
+            shell=True,
+            stdout=stdout,
+            stderr=stderr
+        )
+
+        # Wait for the process to complete
+        returncode = process.wait()
+
+        if fatal and returncode != 0:
+            raise subprocess.CalledProcessError(returncode, script)
+
+        LOG.debug(f"Bash script exited with return code {returncode}")
+
+    except subprocess.CalledProcessError as e:
+        LOG.critical(f"A critical error occurred in a bash script, which exited with code {
+                     e.returncode}.")
+        LOG.critical(f"Script: {script}")
+        # The exception will be caught by the main loop, which will shut down the VM
+        raise HandlerError(f"Bash script failed with exit code {e.returncode}")
+    except FileNotFoundError:
+        raise HandlerError(
+            "Could not find the specified shell or command to execute the script.")
+    except KeyboardInterrupt:
+        LOG.info("Bash script interrupted by user (Ctrl+C)")
+        if process:
+            try:
+                # Immediately send SIGKILL to the entire process group for immediate termination
+                os.killpg(process.pid, signal.SIGKILL)
+                # Wait briefly for the process to die
+                process.wait(timeout=2)
+            except (OSError, subprocess.TimeoutExpired):
+                # Process might already be dead
+                pass
+        raise
+
+
+@StageHandler.method
+def echo(state, text: str):
+    print(text)
+
+
+@StageHandler.method
+def snapshot(state, drive: str, name: str,
+             overwrite: bool = False, **kwargs):
+    if drive not in state.storage:
+        raise HandlerError(f"Drive {drive} not exists")
+
+    state.storage[drive].snapshot(name, overwrite)
+
+
+@StageHandler.method
+def device_add(state, driver: str, id: str, **kwargs):
+    """Adds a device to the VM using QMP."""
+    command = qmp_commands.qmp_device_add(driver, id, **kwargs)
+    qmp_run_command(state.vm, **command)
+
+
+@StageHandler.method
+def device_del(state, id: str):
+    """Removes a device from the VM using QMP."""
+    command = qmp_commands.qmp_device_del(id)
+    qmp_run_command(state.vm, **command)

maqet/ipc/__init__.py

@@ -0,0 +1,29 @@
+"""
+IPC (Inter-Process Communication) Module
+
+Provides Unix domain socket-based IPC for communication between
+CLI processes and VM runner processes.
+
+Components:
+- UnixSocketIPCServer: Server for VM runner processes
+- RunnerClient: Client for CLI processes
+- retry: Retry logic with exponential backoff and circuit breaker
+
+Protocol: JSON-RPC over Unix domain sockets
+"""
+
+from .unix_socket_server import UnixSocketIPCServer
+from .runner_client import RunnerClient
+from .retry import (
+    retry_with_backoff,
+    async_retry_with_backoff,
+    CircuitBreaker,
+)
+
+__all__ = [
+    "UnixSocketIPCServer",
+    "RunnerClient",
+    "retry_with_backoff",
+    "async_retry_with_backoff",
+    "CircuitBreaker",
+]

maqet/ipc/retry.py

@@ -0,0 +1,265 @@
+"""
+IPC Retry Logic with Exponential Backoff
+
+Provides decorators and utilities for retrying IPC operations that may fail
+due to transient network issues, temporary unavailability, or timing issues.
+
+Features:
+- Exponential backoff with configurable base delay
+- Configurable max attempts
+- Selective exception handling (retry only on specific exceptions)
+- Debug logging for retry attempts
+- Synchronous and asynchronous decorator support
+"""
+
+import asyncio
+import functools
+import time
+from typing import Callable, Tuple, Type, Union
+
+from ..logger import LOG
+
+
+def retry_with_backoff(
+    max_attempts: int = 3,
+    backoff_base: float = 0.5,
+    exceptions: Tuple[Type[Exception], ...] = (
+        ConnectionRefusedError,
+        FileNotFoundError,
+    ),
+):
+    """
+    Decorator for synchronous functions with exponential backoff retry logic.
+
+    Retries function on specified exceptions using exponential backoff:
+    - Attempt 1: Immediate
+    - Attempt 2: Wait backoff_base seconds (0.5s default)
+    - Attempt 3: Wait backoff_base * 2 seconds (1s default)
+    - Attempt 4: Wait backoff_base * 4 seconds (2s default)
+    - etc.
+
+    Args:
+        max_attempts: Maximum number of attempts (default: 3)
+        backoff_base: Base delay in seconds for exponential backoff (default: 0.5)
+        exceptions: Tuple of exception types to retry on
+                   (default: ConnectionRefusedError, FileNotFoundError)
+
+    Returns:
+        Decorated function that retries on transient failures
+
+    Example:
+        @retry_with_backoff(max_attempts=3, backoff_base=0.5)
+        def connect_to_service():
+            return service.connect()
+
+    Note:
+        - Only retries on specified exceptions (transient errors)
+        - Other exceptions propagate immediately (permanent errors)
+        - Logs retry attempts at DEBUG level
+        - Final failure logs at WARNING level
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            last_exception = None
+
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    return func(*args, **kwargs)
+
+                except exceptions as e:
+                    last_exception = e
+
+                    if attempt < max_attempts:
+                        # Calculate delay with exponential backoff
+                        delay = backoff_base * (2 ** (attempt - 1))
+                        LOG.debug(
+                            f"IPC retry attempt {attempt}/{max_attempts} failed: {e}. "
+                            f"Retrying in {delay:.2f}s..."
+                        )
+                        time.sleep(delay)
+                    else:
+                        # Max attempts reached
+                        LOG.warning(
+                            f"IPC operation failed after {max_attempts} attempts: {e}"
+                        )
+
+                except Exception as e:
+                    # Non-retryable exception, propagate immediately
+                    LOG.debug(f"Non-retryable exception in IPC operation: {e}")
+                    raise
+
+            # If we get here, all attempts failed
+            raise last_exception
+
+        return wrapper
+
+    return decorator
+
+
+def async_retry_with_backoff(
+    max_attempts: int = 3,
+    backoff_base: float = 0.5,
+    exceptions: Tuple[Type[Exception], ...] = (
+        ConnectionRefusedError,
+        FileNotFoundError,
+        OSError,
+    ),
+):
+    """
+    Decorator for asynchronous functions with exponential backoff retry logic.
+
+    Async version of retry_with_backoff(). Retries async function on specified
+    exceptions using exponential backoff with asyncio.sleep.
+
+    Args:
+        max_attempts: Maximum number of attempts (default: 3)
+        backoff_base: Base delay in seconds for exponential backoff (default: 0.5)
+        exceptions: Tuple of exception types to retry on
+                   (default: ConnectionRefusedError, FileNotFoundError, OSError)
+
+    Returns:
+        Decorated async function that retries on transient failures
+
+    Example:
+        @async_retry_with_backoff(max_attempts=3, backoff_base=0.5)
+        async def connect_to_socket():
+            return await asyncio.open_unix_connection("/path/to/socket")
+
+    Note:
+        - Only retries on specified exceptions (transient errors)
+        - Other exceptions propagate immediately (permanent errors)
+        - Logs retry attempts at DEBUG level
+        - Final failure logs at WARNING level
+        - Uses asyncio.sleep for non-blocking delays
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            last_exception = None
+
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    return await func(*args, **kwargs)
+
+                except exceptions as e:
+                    last_exception = e
+
+                    if attempt < max_attempts:
+                        # Calculate delay with exponential backoff
+                        delay = backoff_base * (2 ** (attempt - 1))
+                        LOG.debug(
+                            f"IPC retry attempt {attempt}/{max_attempts} failed: {e}. "
+                            f"Retrying in {delay:.2f}s..."
+                        )
+                        await asyncio.sleep(delay)
+                    else:
+                        # Max attempts reached
+                        LOG.warning(
+                            f"IPC operation failed after {max_attempts} attempts: {e}"
+                        )
+
+                except Exception as e:
+                    # Non-retryable exception, propagate immediately
+                    LOG.debug(f"Non-retryable exception in IPC operation: {e}")
+                    raise
+
+            # If we get here, all attempts failed
+            raise last_exception
+
+        return wrapper
+
+    return decorator
+
+
+class CircuitBreaker:
+    """
+    Circuit breaker pattern for IPC operations.
+
+    Prevents repeated attempts to connect to a failing service by
+    "opening the circuit" after a threshold of failures.
+
+    States:
+    - CLOSED: Normal operation, requests pass through
+    - OPEN: Too many failures, requests fail immediately
+    - HALF_OPEN: Testing if service recovered, allow one request
+
+    Usage:
+        breaker = CircuitBreaker(failure_threshold=5, timeout=60)
+
+        if breaker.is_open():
+            raise Exception("Circuit breaker open")
+
+        try:
+            result = connect_to_service()
+            breaker.record_success()
+        except Exception:
+            breaker.record_failure()
+            raise
+
+    Note:
+        - Circuit opens after failure_threshold consecutive failures
+        - After timeout seconds, circuit enters HALF_OPEN state
+        - One successful request in HALF_OPEN state closes circuit
+        - Thread-safe for concurrent access
+    """
+
+    def __init__(self, failure_threshold: int = 5, timeout: float = 60.0):
+        """
+        Initialize circuit breaker.
+
+        Args:
+            failure_threshold: Number of failures before opening circuit
+            timeout: Seconds to wait before attempting recovery (HALF_OPEN state)
+        """
+        self.failure_threshold = failure_threshold
+        self.timeout = timeout
+        self.failure_count = 0
+        self.last_failure_time = None
+        self.state = "CLOSED"  # CLOSED, OPEN, HALF_OPEN
+
+    def is_open(self) -> bool:
+        """
+        Check if circuit is open (blocking requests).
+
+        Returns:
+            True if circuit is open and requests should be blocked
+        """
+        if self.state == "OPEN":
+            # Check if timeout elapsed, transition to HALF_OPEN
+            if (
+                self.last_failure_time
+                and time.time() - self.last_failure_time >= self.timeout
+            ):
+                LOG.debug("Circuit breaker entering HALF_OPEN state")
+                self.state = "HALF_OPEN"
+                return False
+            return True
+        return False
+
+    def record_success(self) -> None:
+        """Record successful operation, reset failure count."""
+        if self.state == "HALF_OPEN":
+            LOG.debug("Circuit breaker closing (recovered)")
+        self.failure_count = 0
+        self.state = "CLOSED"
+
+    def record_failure(self) -> None:
+        """Record failed operation, potentially open circuit."""
+        self.failure_count += 1
+        self.last_failure_time = time.time()
+
+        if self.failure_count >= self.failure_threshold:
+            LOG.warning(
+                f"Circuit breaker opening after {self.failure_count} failures"
+            )
+            self.state = "OPEN"
+
+    def reset(self) -> None:
+        """Manually reset circuit breaker to CLOSED state."""
+        LOG.debug("Circuit breaker manually reset")
+        self.failure_count = 0
+        self.last_failure_time = None
+        self.state = "CLOSED"