nest/uvm/firecracker/process.py

"""Firecracker process lifecycle and PID identity handling."""

from __future__ import annotations

import os
import signal
import subprocess
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

from ..config import Settings
from ..domain import VmRecord
from ..errors import FirecrackerError


class _PopenLike(Protocol):
    pid: int

    def poll(self) -> int | None: ...


@dataclass(frozen=True, slots=True)
class ProcessInfo:
    pid: int
    start_time: str


class FirecrackerProcessManager:
    """Start, observe, and terminate detached Firecracker processes."""

    def __init__(
        self,
        settings: Settings,
        *,
        popen: Callable[..., _PopenLike] = subprocess.Popen,
        sleep: Callable[[float], None] = time.sleep,
        monotonic: Callable[[], float] = time.monotonic,
    ) -> None:
        self._settings = settings
        self._popen = popen
        self._sleep = sleep
        self._monotonic = monotonic

    def start(self, vm: VmRecord) -> ProcessInfo:
        socket_path = Path(vm.socket)
        if socket_path.exists():
            try:
                socket_path.unlink()
            except OSError as error:
                raise FirecrackerError(
                    f"could not remove stale Firecracker socket {socket_path}: {error}"
                ) from error

        log_path = Path(vm.log)
        log_path.parent.mkdir(parents=True, exist_ok=True)
        process: _PopenLike | None = None
        try:
            with log_path.open("ab", buffering=0) as log_file:
                process = self._popen(
                    [str(self._settings.firecracker_binary), "--api-sock", str(socket_path)],
                    stdout=log_file,
                    stderr=log_file,
                    start_new_session=True,
                )
            self._wait_for_socket(process, socket_path, log_path)
            start_time = self.process_start_time(process.pid)
            if start_time is None:
                raise FirecrackerError(
                    f"could not establish a safe process identity for Firecracker process {process.pid}"
                )
            return ProcessInfo(pid=process.pid, start_time=start_time)
        except OSError as error:
            if process is not None:
                self._cleanup_started_process(process, socket_path)
            raise FirecrackerError(f"could not start Firecracker: {error}") from error
        except BaseException:
            if process is not None:
                self._cleanup_started_process(process, socket_path)
            raise

    def is_alive(self, vm: VmRecord) -> bool:
        if vm.pid is None or vm.pid <= 0 or vm.process_start_time is None:
            return False
        try:
            os.kill(vm.pid, 0)
        except ProcessLookupError:
            return False
        except PermissionError:
            return True

        return self.process_start_time(vm.pid) == vm.process_start_time

    def terminate(self, vm: VmRecord) -> None:
        if vm.pid is None:
            self._remove_socket(Path(vm.socket))
            return
        if vm.pid <= 0:
            raise FirecrackerError(f"invalid persisted Firecracker PID: {vm.pid}")
        if vm.process_start_time is None:
            if not self._pid_exists(vm.pid):
                self._remove_socket(Path(vm.socket))
                return
            raise FirecrackerError(
                f"refusing to signal Firecracker PID {vm.pid} without a process identity token"
            )
        if not self.is_alive(vm):
            self._remove_socket(Path(vm.socket))
            return

        try:
            os.kill(vm.pid, signal.SIGTERM)
        except ProcessLookupError:
            self._remove_socket(Path(vm.socket))
            return
        except OSError as error:
            raise FirecrackerError(f"could not stop Firecracker process {vm.pid}: {error}") from error

        deadline = self._monotonic() + self._settings.terminate_timeout_s
        while self.is_alive(vm) and self._monotonic() < deadline:
            self._sleep(0.05)
        if self.is_alive(vm):
            self._terminate_process_id(vm.pid, force=True)
        self._remove_socket(Path(vm.socket))

    def process_start_time(self, pid: int) -> str | None:
        """Read Linux proc start time so a reused PID is not mistaken for a VM."""

        try:
            stat = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
        except OSError:
            return None
        closing_parenthesis = stat.rfind(")")
        if closing_parenthesis < 0:
            return None
        fields = stat[closing_parenthesis + 2 :].split()
        try:
            return fields[19]
        except IndexError:
            return None

    def _wait_for_socket(
        self,
        process: _PopenLike,
        socket_path: Path,
        log_path: Path,
    ) -> None:
        deadline = self._monotonic() + self._settings.api_socket_timeout_s
        while self._monotonic() < deadline:
            if socket_path.exists():
                return
            if process.poll() is not None:
                raise FirecrackerError(f"Firecracker exited early; see {log_path}")
            self._sleep(0.05)
        raise FirecrackerError("timed out waiting for Firecracker API socket")

    @staticmethod
    def _remove_socket(socket_path: Path) -> None:
        try:
            socket_path.unlink()
        except FileNotFoundError:
            pass
        except OSError:
            # A stale socket is cleaned on the next launch; never hide a successful stop.
            pass

    @staticmethod
    def _pid_exists(pid: int) -> bool:
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            return False
        except PermissionError:
            return True
        return True

    def _cleanup_started_process(self, process: _PopenLike, socket_path: Path) -> None:
        try:
            self._terminate_process_id(process.pid)
        except FirecrackerError:
            pass
        deadline = self._monotonic() + self._settings.terminate_timeout_s
        while process.poll() is None and self._monotonic() < deadline:
            self._sleep(0.05)
        if process.poll() is None:
            try:
                self._terminate_process_id(process.pid, force=True)
            except FirecrackerError:
                pass
        self._remove_socket(socket_path)

    @staticmethod
    def _terminate_process_id(pid: int, *, force: bool = False) -> None:
        if pid <= 0:
            raise FirecrackerError(f"invalid Firecracker PID: {pid}")
        signal_to_send = signal.SIGKILL if force else signal.SIGTERM
        try:
            os.kill(pid, signal_to_send)
        except ProcessLookupError:
            pass
        except OSError as error:
            action = "force-stop" if force else "stop"
            raise FirecrackerError(f"could not {action} Firecracker process {pid}: {error}") from error