scripts/tests/chiptest/concurrency/process.py

# Copyright (c) 2026 Project CHIP Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
import dataclasses
import enum
import functools
import logging
import multiprocessing
import os
import signal
import threading
import traceback
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterator
from multiprocessing.context import SpawnContext
from multiprocessing.managers import SyncManager, ValueProxy
from types import TracebackType
from typing import ClassVar, Concatenate, Generic, ParamSpec, Protocol, Self, TypeVar

from chiptest.concurrency.work_queue import CancellableQueue, EndOfQueue, QueueCancelled, wait_for_mp_managed
from chiptest.log_config import LogConfig

log = logging.getLogger(__name__)


class WithName(Protocol):
    @property
    def name(self) -> str: ...


S = TypeVar("S", bound="WithName")
P = ParamSpec("P")
R = TypeVar("R")


def with_annotated_exception(fn: Callable[Concatenate[S, P], R]) -> Callable[Concatenate[S, P], R]:
    """Decorator to enrich exceptions from thread and process methods with their names for easier debugging."""

    @functools.wraps(fn)
    def wrapper(self: S, *args: P.args, **kwargs: P.kwargs) -> R:
        try:
            return fn(self, *args, **kwargs)
        except BaseException as e:
            kind = ("thread" if isinstance(self, threading.Thread) else
                    "process" if isinstance(self, (multiprocessing.Process, WrappedProcess)) else
                    self.__class__.__name__)

            e.add_note(f"Exception in {kind} {self.name}")
            raise

    return wrapper


@dataclasses.dataclass
class ProcessConfig:
    """Process configuration template."""

    DEFAULT_START_TIMEOUT_SEC: ClassVar[float] = 4.0
    DEFAULT_STOP_TIMEOUT_SEC: ClassVar[float] = 4.0
    DEFAULT_TERMINATION_TIMEOUT_SEC: ClassVar[float] = 2.0

    id: int = 0
    name: str = "Process{id}"
    log_config: LogConfig = dataclasses.field(default_factory=LogConfig)
    start_timeout_sec: float = DEFAULT_START_TIMEOUT_SEC
    stop_timeout_sec: float = DEFAULT_STOP_TIMEOUT_SEC
    termination_timeout_sec: float = DEFAULT_TERMINATION_TIMEOUT_SEC

    def with_formatted_name(self) -> Self:
        """Format the name using the id and process hierarchy. Return a new instance with the formatted name."""
        name = self.name.format(id=self.id)

        # Indicate process hierarchy in the name so that it's easier to understand the dependencies in the logs.
        if (parent_proc_name := multiprocessing.current_process().name) != "MainProcess":
            name = f"{parent_proc_name}/{name}"

        return dataclasses.replace(self, name=name)


class ProcessPhase(enum.IntEnum):
    """Phase of the process lifecycle."""

    NOT_STARTED = enum.auto()
    """The process has not started yet."""

    UNINITIALIZED = enum.auto()
    """The process has started but has not finished initialization yet."""

    READY = enum.auto()
    """The process has finished initialization and is ready to work. Can cycle between READY and WORKING phases multiple times."""

    WORKING = enum.auto()
    """The process is working on a task. Optional and can be used to distinguish between idle and busy processes."""

    CLOSED = enum.auto()
    """The process has finished all work and is closed (exited). Mind that it could have been closed due to an exception."""


class ProcessState:
    """
    Shared process state for synchronized lifecycle coordination and error propagation between the parent and child processes.

    Can be used as a context manager to synchronize access to the state. Needed only if accessing multiple fields atomically,
    otherwise individual fields are synchronized internally.
    """

    def __init__(self, mp_manager: SyncManager, config: ProcessConfig) -> None:
        self._config = config
        self._phase: ValueProxy[ProcessPhase] = mp_manager.Value(object, ProcessPhase.NOT_STARTED)
        self._exception: ValueProxy[BaseException | None] = mp_manager.Value(object, None)
        self._state_changed = mp_manager.Condition()

    @property
    def phase(self) -> ProcessPhase:
        return self._phase.get()

    @phase.setter
    def phase(self, value: ProcessPhase) -> None:
        with self._state_changed:
            self._phase.set(value)
            self._state_changed.notify_all()

    @property
    def exception(self) -> BaseException | None:
        return self._exception.get()

    @exception.setter
    def exception(self, value: BaseException | None) -> None:
        with self._state_changed:
            if isinstance(value, BaseException):
                value.add_note(f"Exception in process {self._config.name}")

                # Traceback is not automatically propagated across process boundaries, so we need to add it to the exception
                # manually as a note.
                value.add_note("".join(traceback.format_exception(value)))
            self._exception.set(value)
            self._state_changed.notify_all()

    def __enter__(self) -> bool:
        return self._state_changed.__enter__()

    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None):
        return self._state_changed.__exit__(exc_type, exc_val, exc_tb)

    def wait_for(self, predicate: Callable[[ProcessPhase, BaseException | None], bool], timeout: float | None = None):
        """
        Wait until the predicate returns True for the current phase and/or exception.

        Timeout as for `Condition.wait_for()`, i.e. it can be a positive float for maximum wait time, or None to wait indefinitely.

        Returns the result of the predicate.

        We need the wait_for_mp_managed wrapper because otherwise we wouldn't be able to catch a KeyboardInterrupt for the condition
        which is managed by multiprocessing.Manager, as the manager process explicitly ignores SIGINT.
        """
        with self._state_changed:
            return wait_for_mp_managed(self._state_changed, lambda: predicate(self._phase.get(), self._exception.get()), timeout)

    @contextlib.contextmanager
    def working_context(self) -> Iterator[None]:
        """Context manager to set the process phase to WORKING for the duration of the context."""
        self.phase = ProcessPhase.WORKING
        try:
            yield
        finally:
            self.phase = ProcessPhase.READY


WorkerConfigT = TypeVar("WorkerConfigT", bound=ProcessConfig)
WorkRequestT = TypeVar("WorkRequestT")
WorkResponseT = TypeVar("WorkResponseT")


class WrappedProcess(ABC, Generic[WorkerConfigT, WorkRequestT, WorkResponseT]):
    """
    Base class for wrapped Python subprocesses.

    WrappedProcess owns a `multiprocessing.Process` and coordinates its lifecycle through shared `ProcessState` and work/response
    queues.

    Intended wrapping flow:
    1. The caller enters `mp_wrapped_spawn_context()` from `concurrent/context.py`.
    2. The context always uses the spawn start method, and on Linux can optionally install a temporary wrapper executable that
       launches Python through `wrapper_linux`.
    3. That `SpawnContext` is passed to WrappedProcess and used to create the child process.
    4. Parent-side lifecycle is driven by `start()` / `stop()`.
    5. Child-side lifecycle is driven by `run()` and the `_proc_*` hook methods.
    6. On context exit, `mp_wrapped_spawn_context()` restores the original multiprocessing executable and removes the wrapper.

    The lifecycle of the subprocess is coordinated through `ProcessState` and follows these phases:
    1. Parent calls `start()`.
    2. Subprocess enters `run()` and sets phase `UNINITIALIZED`.
    3. `_proc_init()` is invoked once for subclass-specific initialization. All
       resources that require cleanup should be registered in the provided `ExitStack`.
    4. On success, phase becomes `READY`.
    5. `_proc_work()` executes the main loop (by default waits for cancellation). Implementations may optionally toggle between
       `READY` and `WORKING` to expose idle vs active periods. You can use working_context() for that.
    6. After finishing work, the resources initialized in `_proc_init()` are cleaned up by the exit stack.
    7. Phase becomes `CLOSED` when the subprocess exits.

    Error handling behavior:
    - Unhandled exceptions from subprocess code are stored in `state.exception`.
    - `start()` waits for initialization and fails if the process reports an exception, closes early, or times out.
    - `stop()` attempts graceful shutdown first, then escalates with SIGINT, SIGTERM, and SIGKILL with configured timeouts.

    Subclasses are expected to implement `_proc_init()`, and may optionally override `_proc_work()`.
    """
    # Methods run in the parent process.

    def __init__(self, mp_context: SpawnContext, mp_manager: SyncManager, config: WorkerConfigT,
                 work_queue: CancellableQueue[WorkRequestT], rsp_queue: CancellableQueue[WorkResponseT]) -> None:
        # Neither mp_context or mp_manager should be saved as fields, as they are not picklable between processes. They can be used
        # to initialize some shared resources in the constructor.

        self._config = config
        self._work_queue = work_queue
        self._rsp_queue = rsp_queue
        self.state = ProcessState(mp_manager, config)

        # Create multiprocessing.Process in the given context.
        self._proc = mp_context.Process(target=self.run, name=self._config.name)
        self._stopped = False

    def __getstate__(self) -> dict[str, object]:
        state = self.__dict__.copy()
        # Drop _proc field to exclude it from pickling on process spawn.
        state.pop("_proc", None)
        return state

    @property
    def name(self) -> str:
        return self._config.name

    @with_annotated_exception
    def start(self) -> None:
        """
        Start the subprocess and wait until initialization completes.

        After triggering the process start, wait until phase leaves `NOT_STARTED`/`UNINITIALIZED` and then validate final startup
        state: `READY` is success, and `CLOSED` or any recorded exception is failure.
        """
        if self._stopped:
            raise RuntimeError("Cannot start a process that has been already stopped")

        log.debug("Starting process %s", self.name)
        self._proc.start()

        try:
            if not self.state.wait_for(lambda phase, _: phase not in (ProcessPhase.NOT_STARTED, ProcessPhase.UNINITIALIZED),
                                       self._config.start_timeout_sec):
                raise TimeoutError("Timeout when waiting for initialization")

            with self.state:
                # Propagate KeyboardInterrupt as is.
                if isinstance(self.state.exception, KeyboardInterrupt):
                    raise self.state.exception

                if self.state.exception is not None:
                    raise RuntimeError("Process failed initialization") from self.state.exception

                match self.state.phase:
                    case ProcessPhase.READY:
                        log.debug("Process %s started successfully", self.name)
                    case ProcessPhase.CLOSED:
                        raise RuntimeError("Process closed immediately after initialization")
                    case _:
                        raise RuntimeError(f"Process is in an unexpected state: {self.state.phase!r}")
        except BaseException as start_exc:
            if not isinstance(start_exc, KeyboardInterrupt):
                log.error("Stopping process %s on failure during initialization", self.name)
            try:
                self.stop()
            except KeyboardInterrupt:
                raise
            except BaseException as stop_exc:
                log.error("Error when stopping process %s after failure during initialization: %r", self.name, stop_exc)
                raise stop_exc.with_traceback(stop_exc.__traceback__) from start_exc
            raise

    def has_stopped(self, timeout: float) -> bool:
        self._proc.join(timeout)
        if not self._proc.is_alive():
            log.debug("Process %s got stopped", self.name)
            return True
        return False

    @with_annotated_exception
    def stop(self) -> None:
        """
        Stop the subprocess with escalating termination signals.

        The method waits for a graceful exit first (typically triggered when the work queue is cancelled by its owner). If still
        alive, it escalates to SIGINT, then SIGTERM, then SIGKILL.
        """
        if self._stopped:
            log.debug("Process %s is already stopped", self.name)
            return

        try:
            # First check if the process is alive before trying to join, as it is an error to attempt to join a process before it
            # has been started.
            if not self._proc.is_alive():
                log.debug("Process %s hasn't been started yet or has been already stopped", self.name)
                return

            # Wait for the external work queue to be cancelled by its owner, which should signal the process to gracefully stop.
            if self.has_stopped(self._config.stop_timeout_sec):
                return

            log.debug("Sending interrupt signal to process %s", self.name)
            if self._proc.pid is not None:
                os.kill(self._proc.pid, signal.SIGINT)  # TODO Python 3.14: self._proc.interrupt()
            if self.has_stopped(self._config.stop_timeout_sec):
                return

            log.warning("Timeout when waiting for process %s to stop. Sending terminate signal", self.name)
            self._proc.terminate()
            if self.has_stopped(self._config.termination_timeout_sec):
                return

            log.warning("Timeout when waiting for process %s to terminate. Sending kill signal", self.name)
            self._proc.kill()
            if self.has_stopped(self._config.termination_timeout_sec):
                return

            raise TimeoutError(f"Failed to terminate the process {self.name}. May become a zombie")
        finally:
            self._stopped = True
            if isinstance(self.state.exception, KeyboardInterrupt):
                raise self.state.exception
            if self.state.exception is not None:
                raise RuntimeError("Process reported an exception during execution") from self.state.exception

    # Methods run in the subprocess.

    def run(self) -> None:
        """
        Subprocess entrypoint coordinating init/work/cleanup and state.

        This method drives the lifecycle hooks in order: `_proc_init()` -> `_proc_work()`. It updates `state.phase` and captures
        unexpected exceptions in `state.exception`.

        We use several nested try/except blocks:
        1. The outermost block captures any exception that passed through other filters, ensures that it is stored in state and
           ensures that the process phase is set to CLOSED at the end.
        2. The middle block filters out expected exceptions related to normal process shutdown (like `QueueCancelled`), so that they
           are not treated as errors in scope of the exit stack.
        3. The innermost blocks around `_proc_init()` and `_proc_work()` enrich the exceptions with additional context about the
           failure stage.
        """
        try:
            # Signal that the process has started initialization.
            self.state.phase = ProcessPhase.UNINITIALIZED

            # Initialize global logger in the subprocess.
            self._config.log_config.set_fmt()

            with contextlib.ExitStack() as stack:
                try:
                    try:
                        log.debug("Initializing")
                        self._proc_init(stack)
                        self.state.phase = ProcessPhase.READY
                        log.debug("Initialized successfully")
                    except BaseException as e:
                        e.add_note("Failure during process initialization")
                        raise

                    try:
                        self._proc_work()
                    except BaseException as e:
                        e.add_note("Failure during process work")
                        raise

                # Capture exceptions that are not errors.
                except QueueCancelled:
                    log.warning("Received a cancel event")
                except EndOfQueue:
                    log.debug("Received end of work signal")
        except BaseException as e:
            if isinstance(e, KeyboardInterrupt):
                log.debug("Process interrupted by user")
            else:
                log.error("Process failed with an exception: %r", e)
            self.state.exception = e
        finally:
            self.state.phase = ProcessPhase.CLOSED

    @abstractmethod
    def _proc_init(self, exit_stack: contextlib.ExitStack) -> None:
        """Initialize subprocess resources before work begins.

        The long-lived resources that need cleanup should be registered in the provided `ExitStack`.
        """

    def _proc_work(self) -> None:
        """
        Perform the work.

        Default behavior waits for cancellation. Override to consume `self._work_queue` and publish activity by toggling
        `state.phase` between `READY` and `WORKING` as appropriate.
        """
        self._work_queue.wait_for_cancelled()