Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
Toggle table of contents sidebar

# Source code for pytestshellutils.utils.processes

    # Copyright 2021-2024 VMware, Inc.
    # SPDX-License-Identifier: Apache-2.0
    """
    Process related utilities.
    """
    import errno
    import json
    import logging
    import pprint
    import signal
    import weakref
    from typing import Any
    from typing import Dict
    from typing import List
    from typing import Optional

    import attr
    import psutil

    try:
        from pytest import LineMatcher
    except ImportError:  # pragma: no cover
        # Older pytest
        from _pytest.pytester import LineMatcher

    from pytestshellutils.utils import warn_until

    log = logging.getLogger(__name__)

    class MatchString(str):
        """
        Simple subclass around ``str`` which provides a ``.matcher`` property.

        This ``.matcher`` property is an instance of :py:class:`~pytest.LineMatcher`
        """

        @property
        def matcher(self) -> LineMatcher:
            """
            Return an instance of :py:class:`~pytest.LineMatcher`.
            """
            return LineMatcher(self.splitlines())

    def convert_string_to_match_string(value: Optional[str]) -> Optional[MatchString]:
        """
        Convert strings into ``MatchString`` instances.
        """
        if isinstance(value, str):
            return MatchString(value)
        return value

    @attr.s(frozen=True, kw_only=True)
    class ProcessResult:
        """
        Wrapper class around a subprocess result.

        This class serves the purpose of having a common result class which will hold the
        resulting data from a subprocess command.

        Keyword Arguments:
            returncode:
                The returncode returned by the process
            stdout:
                The ``stdout`` returned by the process
            stderr:
                The ``stderr`` returned by the process
            cmdline:
                The command line used to start the process
            data:
                The data returned by parsing ``stdout``, when possible.
            data_key:
                When ``stdout`` can be parsed as JSON, sometimes there's a top level key which is not
                that interesting. By using ``data_key``, we define that we're actually only interested
                on the data structure which is keyed by ``data_key``.

        Note:
            Cast :py:class:`~pytestshellutils.utils.processes.ProcessResult` to a string to pretty-print it.
        """

        returncode: int = attr.ib()
        stdout: MatchString = attr.ib(converter=convert_string_to_match_string)
        stderr: MatchString = attr.ib(converter=convert_string_to_match_string)
        cmdline: Optional[List[str]] = attr.ib(default=None)
        data_key: Optional[str] = attr.ib(default=None)
        data: Optional[Dict[Any, Any]] = attr.ib()

        @returncode.validator
        def _validate_returncode(self, attribute: Any, value: int) -> None:
            """
            Validate the value type.
            """
            if not isinstance(value, int):
                raise ValueError(f"'returncode' needs to be an integer, not '{type(value)}'")

        @data.default
        def _default_data(self) -> Optional[Dict[Any, Any]]:
            """
            Try to parse the passed ``stdout`` as JSON as the default data value.
            """
            stdout: Optional[str] = self.stdout.strip() if self.stdout else None
            if stdout:
                try:
                    data: Optional[Dict[Any, Any]] = json.loads(stdout.strip())
                    if data and self.data_key and self.data_key in data:
                        data = data[self.data_key]
                    return data
                except ValueError:
                    pass
            return None

        @property
        def exitcode(self) -> int:  # pragma: no cover
            """
            Return the process returncode.

            This property is deprecated and should not be used.
            It only exists to support projects that are migrating from
            pytest-salt-factories versions. Use ``.returncode`` instead.
            """
            warn_until(
                "2.0.0",
                "The '.exitcode' property is deprecated and will cease to exist after "
                "pytest-shell-utilities {version}. Please use '.returncode' instead.",
            )
            return self.returncode

        @property
        def json(self) -> Optional[Dict[Any, Any]]:  # pragma: no cover
            """
            Return the process output parsed as JSON, if possible.

            This property is deprecated and should not be used.
            It only exists to support projects that are migrating from
            pytest-salt-factories versions. Use ``.data`` instead.
            """
            warn_until(
                "2.0.0",
                "The '.json' property is deprecated and will cease to exist after "
                "pytest-shell-utilities {version}. Please use '.data' instead.",
            )
            return self.data

        def __str__(self) -> str:
            """
            String representation of the class.
            """
            message = self.__class__.__name__
            if self.cmdline:
                message += f"\n Command Line: {self.cmdline}"
            if self.returncode is not None:
                message += f"\n Returncode: {self.returncode}"
            if (self.stdout and self.stdout.strip()) or (self.stderr and self.stderr.strip()):
                message += "\n Process Output:"
            if self.stdout and self.stdout.strip():
                message += f"\n   >>>>> STDOUT >>>>>\n{self.stdout}\n   <<<<< STDOUT <<<<<"
            if self.stderr and self.stderr.strip():
                message += f"\n   >>>>> STDERR >>>>>\n{self.stderr}\n   <<<<< STDERR <<<<<"
            if self.data:
                message += f"\n Parsed JSON Data:\n{self._to_printable_data(self.data)}"
            return message + "\n"

        @staticmethod
        def _to_printable_data(data: Dict[str, Any]) -> str:
            """
            Convert a data dictionary into a JSON string.
            """
            return "\n".join(
                f"   {line}" for line in json.dumps(data, indent=2, sort_keys=False).splitlines()
            ).rstrip()

    def collect_child_processes(pid: int) -> List[psutil.Process]:
        """
        Try to collect any started child processes of the provided pid.

        Arguments:
            pid:
                The PID of the process

        Returns:
            List of child processes
        """
        # Let's get the child processes of the started subprocess
        children: List[psutil.Process]
        try:
            parent = psutil.Process(pid)
            children = parent.children(recursive=True)
        except psutil.NoSuchProcess:
            children = []
        return children

    def _get_cmdline(proc: psutil.Process) -> Optional[Any]:
        # pylint: disable=protected-access
        try:
            return proc._cmdline
        except AttributeError:
            # Cache the cmdline since that will be inaccessible once the process is terminated
            # and we use it in log calls
            try:
                cmdline = proc.cmdline()
            except (psutil.NoSuchProcess, psutil.AccessDenied):  # pragma: no cover
                # OSX is more restrictive about the above information
                cmdline = None
            except OSError:  # pragma: no cover
                # On Windows we've seen something like:
                #   File " c: ... \lib\site-packages\pytestsalt\utils\__init__.py", line 182, in terminate_process
                #     terminate_process_list(process_list, kill=slow_stop is False, slow_stop=slow_stop)
                #   File " c: ... \lib\site-packages\pytestsalt\utils\__init__.py", line 130, in terminate_process_list
                #     _terminate_process_list(process_list, kill=kill, slow_stop=slow_stop)
                #   File " c: ... \lib\site-packages\pytestsalt\utils\__init__.py", line 78, in _terminate_process_list
                #     cmdline = process.cmdline()
                #   File " c: ... \lib\site-packages\psutil\__init__.py", line 786, in cmdline
                #     return self._proc.cmdline()
                #   File " c: ... \lib\site-packages\psutil\_pswindows.py", line 667, in wrapper
                #     return fun(self, *args, **kwargs)
                #   File " c: ... \lib\site-packages\psutil\_pswindows.py", line 745, in cmdline
                #     ret = cext.proc_cmdline(self.pid, use_peb=True)
                #   OSError: [WinError 299] Only part of a ReadProcessMemory or WriteProcessMemory request was completed: 'originated from ReadProcessMemory(ProcessParameters)
                cmdline = None
            except RuntimeError:  # pragma: no cover
                # Also on windows
                # saltfactories\utils\processes\helpers.py:68: in _get_cmdline
                #     cmdline = proc.as_dict()
                # c: ... \lib\site-packages\psutil\__init__.py:634: in as_dict
                #     ret = meth()
                # c: ... \lib\site-packages\psutil\__init__.py:1186: in memory_full_info
                #     return self._proc.memory_full_info()
                # c: ... \lib\site-packages\psutil\_pswindows.py:667: in wrapper
                #     return fun(self, *args, **kwargs)
                # _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
                #
                # self = <psutil._pswindows.Process object at 0x0000029B7FDA5558>
                #
                #     @wrap_exceptions
                #     def memory_full_info(self):
                #         basic_mem = self.memory_info()
                # >       uss = cext.proc_memory_uss(self.pid)
                # E       RuntimeError: NtQueryVirtualMemory failed
                #
                # c: ... \lib\site-packages\psutil\_pswindows.py:806: RuntimeError
                cmdline = None

            if not cmdline:  # pragma: no cover
                try:
                    cmdline = proc.as_dict()
                except psutil.NoSuchProcess:
                    cmdline = f"<could not be retrived; dead process: {proc}>"
                except (psutil.AccessDenied, OSError):  # pragma: no cover
                    cmdline = weakref.proxy(proc)
            proc._cmdline = cmdline
        return proc._cmdline
        # pylint: enable=protected-access

    def _terminate_process_list(
        process_list: List[psutil.Process], kill: bool = False, slow_stop: bool = False
    ) -> None:
        log.info(
            "Terminating process list:\n%s",
            pprint.pformat([_get_cmdline(proc) for proc in process_list]),
        )
        for process in process_list[:]:  # Iterate over copy of the list
            if not psutil.pid_exists(process.pid):
                process_list.remove(process)
                continue
            try:
                if not kill and process.status() == psutil.STATUS_ZOMBIE:  # pragma: no cover
                    # Zombie processes will exit once child processes also exit
                    continue
                if kill:
                    log.info("Killing process(%s): %s", process.pid, _get_cmdline(process))
                    process.kill()
                else:
                    log.info("Terminating process(%s): %s", process.pid, _get_cmdline(process))
                    try:
                        if slow_stop:
                            # Allow coverage data to be written down to disk
                            process.send_signal(signal.SIGTERM)
                            try:
                                process.wait(2)
                            except psutil.TimeoutExpired:  # pragma: no cover
                                if psutil.pid_exists(process.pid):
                                    continue
                        else:
                            process.terminate()
                    except OSError as exc:  # pragma: no cover
                        if exc.errno not in (errno.ESRCH, errno.EACCES):
                            raise
                if not psutil.pid_exists(process.pid):
                    process_list.remove(process)
            except psutil.NoSuchProcess:
                process_list.remove(process)

    def terminate_process_list(
        process_list: List[psutil.Process], kill: bool = False, slow_stop: bool = False
    ) -> None:
        """
        Terminate a list of processes.

        Arguments:
            process_list:
                An iterable of :py:class:`psutil.Process` instances to terminate

        Keyword Arguments:
            kill:
                Kill the process instead of terminating it.
            slow_stop:
                First try to terminate each process in the list, and if termination was not successful, kill it.

        Returns:
            Nothing.
        """

        def on_process_terminated(proc: psutil.Process) -> None:
            log.info(
                "Process %s terminated with exit code: %s",
                getattr(proc, "_cmdline", proc),
                proc.returncode,
            )

        # Try to terminate processes with the provided kill and slow_stop parameters
        log.info("Terminating process list. 1st step. kill: %s, slow stop: %s", kill, slow_stop)

        # Remove duplicates from the process list
        seen_pids = []
        start_count = len(process_list)
        for proc in process_list[:]:
            if proc.pid in seen_pids:
                process_list.remove(proc)
            seen_pids.append(proc.pid)
        end_count = len(process_list)
        if end_count < start_count:
            log.debug("Removed %d duplicates from the initial process list", start_count - end_count)

        _terminate_process_list(process_list, kill=kill, slow_stop=slow_stop)
        psutil.wait_procs(process_list, timeout=5, callback=on_process_terminated)

        if process_list:
            # If there's still processes to be terminated, retry and kill them if slow_stop is False
            log.info(
                "Terminating process list. 2nd step. kill: %s, slow stop: %s",
                slow_stop is False,
                slow_stop,
            )
            _terminate_process_list(process_list, kill=slow_stop is False, slow_stop=slow_stop)
            psutil.wait_procs(process_list, timeout=5, callback=on_process_terminated)

        if process_list:
            # If there's still processes to be terminated, just kill them, no slow stopping now
            log.info("Terminating process list. 3rd step. kill: True, slow stop: False")
            _terminate_process_list(process_list, kill=True, slow_stop=False)
            psutil.wait_procs(process_list, timeout=5, callback=on_process_terminated)

        if process_list:
            # In there's still processes to be terminated, log a warning about it
            log.warning("Some processes failed to properly terminate: %s", process_list)

    def terminate_process(
        pid: Optional[int] = None,
        process: Optional[psutil.Process] = None,
        children: Optional[List[psutil.Process]] = None,
        kill_children: Optional[bool] = None,
        slow_stop: bool = False,
    ) -> None:
        """
        Try to terminate/kill the started process.

        Keyword Arguments:
            pid:
                The PID of the process
            process:
                An instance of :py:class:`psutil.Process`
            children:
                An iterable of :py:class:`psutil.Process` instances, children to the process being terminated
            kill_children:
                Also try to terminate/kill child processes
            slow_stop:
                First try to terminate each process in the list, and if termination was not successful, kill it.
        """
        children = children or []
        process_list: List[psutil.Process] = []

        if kill_children is None:
            # Always kill children if kill the parent process and kill_children was not set
            kill_children = True if slow_stop is False else kill_children

        if pid and not process:
            try:
                process = psutil.Process(pid)
                process_list.append(process)
            except psutil.NoSuchProcess:
                # Process is already gone
                process = None

        if kill_children:
            if process:
                children.extend(collect_child_processes(process.pid))
            if children:
                process_list.extend(children)

        if process_list:
            if process:
                log.info("Stopping process %s and respective children: %s", process, children)
            else:
                log.info("Terminating process list: %s", process_list)
            terminate_process_list(process_list, kill=slow_stop is False, slow_stop=slow_stop)