diff --git a/BUILD.gn b/BUILD.gn index a2b753239e..fb7b2025f5 100644 --- a/BUILD.gn +++ b/BUILD.gn @@ -266,6 +266,7 @@ group("action_tests") { group("integration_tests") { _default_tc = _default_toolchain_prefix + pw_DEFAULT_C_OPTIMIZATION_LEVEL deps = [ + "$dir_pw_cli/py:process_integration_test.action($_default_tc)", "$dir_pw_rpc:cpp_client_server_integration_test($_default_tc)", "$dir_pw_rpc/py:python_client_cpp_server_test.action($_default_tc)", "$dir_pw_unit_test/py:rpc_service_test.action($_default_tc)", diff --git a/pw_cli/py/BUILD.bazel b/pw_cli/py/BUILD.bazel index e95742249e..c77ce2ca27 100644 --- a/pw_cli/py/BUILD.bazel +++ b/pw_cli/py/BUILD.bazel @@ -51,10 +51,10 @@ py_binary( ) py_test( - name = "plugins_test", + name = "envparse_test", size = "small", srcs = [ - "plugins_test.py", + "envparse_test.py", ], deps = [ ":pw_cli", @@ -62,10 +62,10 @@ py_test( ) py_test( - name = "envparse_test", + name = "plugins_test", size = "small", srcs = [ - "envparse_test.py", + "plugins_test.py", ], deps = [ ":pw_cli", diff --git a/pw_cli/py/BUILD.gn b/pw_cli/py/BUILD.gn index 052ffc6acd..4bb4cd2e38 100644 --- a/pw_cli/py/BUILD.gn +++ b/pw_cli/py/BUILD.gn @@ -46,3 +46,15 @@ pw_python_package("py") { pylintrc = "$dir_pigweed/.pylintrc" mypy_ini = "$dir_pigweed/.mypy.ini" } + +pw_python_script("process_integration_test") { + sources = [ "process_integration_test.py" ] + python_deps = [ ":py" ] + + pylintrc = "$dir_pigweed/.pylintrc" + mypy_ini = "$dir_pigweed/.mypy.ini" + + action = { + stamp = true + } +} diff --git a/pw_cli/py/process_integration_test.py b/pw_cli/py/process_integration_test.py new file mode 100644 index 0000000000..3d05b1a1bf --- /dev/null +++ b/pw_cli/py/process_integration_test.py @@ -0,0 +1,73 @@ +# Copyright 2023 The Pigweed 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 +# +# https://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. +"""Tests for pw_cli.process.""" + +import unittest +import sys +import textwrap + +from pw_cli import process + +import psutil # type: ignore + + +FAST_TIMEOUT_SECONDS = 0.1 +KILL_SIGNALS = set({-9, 137}) +PYTHON = sys.executable + + +class RunTest(unittest.TestCase): + """Tests for process.run.""" + + def test_returns_output(self) -> None: + echo_str = 'foobar' + print_str = f'print("{echo_str}")' + result = process.run(PYTHON, '-c', print_str) + self.assertEqual(result.output, b'foobar\n') + + def test_timeout_kills_process(self) -> None: + sleep_100_seconds = 'import time; time.sleep(100)' + result = process.run( + PYTHON, '-c', sleep_100_seconds, timeout=FAST_TIMEOUT_SECONDS + ) + self.assertIn(result.returncode, KILL_SIGNALS) + + def test_timeout_kills_subprocess(self) -> None: + # Spawn a subprocess which waits for 100 seconds, print its pid, + # then wait for 100 seconds. + sleep_in_subprocess = textwrap.dedent( + f""" + import subprocess + import time + + child = subprocess.Popen( + ['{PYTHON}', '-c', 'import time; print("booh"); time.sleep(100)'] + ) + print(child.pid, flush=True) + time.sleep(100) + """ + ) + result = process.run( + PYTHON, '-c', sleep_in_subprocess, timeout=FAST_TIMEOUT_SECONDS + ) + self.assertIn(result.returncode, KILL_SIGNALS) + # THe first line of the output is the PID of the child sleep process. + child_pid_str, sep, remainder = result.output.partition(b'\n') + del sep, remainder + child_pid = int(child_pid_str) + self.assertFalse(psutil.pid_exists(child_pid)) + + +if __name__ == '__main__': + unittest.main() diff --git a/pw_cli/py/pw_cli/process.py b/pw_cli/py/pw_cli/process.py index e683db8b46..98852e35f0 100644 --- a/pw_cli/py/pw_cli/process.py +++ b/pw_cli/py/pw_cli/process.py @@ -23,6 +23,9 @@ import pw_cli.color import pw_cli.log +import psutil # type: ignore + + _COLOR = pw_cli.color.colors() _LOG = logging.getLogger(__name__) @@ -32,7 +35,12 @@ class CompletedProcess: - """Information about a process executed in run_async.""" + """Information about a process executed in run_async. + + Attributes: + pid: The process identifier. + returncode: The return code of the process. + """ def __init__( self, @@ -84,6 +92,56 @@ async def _run_and_log(program: str, args: Tuple[str, ...], env: dict): return process, bytes(output) +async def _kill_process_and_children( + process: 'asyncio.subprocess.Process', +) -> None: + """Kills child processes of a process with PID `pid`.""" + # Look up child processes before sending the kill signal to the parent, + # as otherwise the child lookup would fail. + pid = process.pid + try: + process_util = psutil.Process(pid=pid) + kill_list = list(process_util.children(recursive=True)) + except psutil.NoSuchProcess: + # Creating the kill list raced with parent process completion. + # + # Don't bother cleaning up child processes of parent processes that + # exited on their own. + kill_list = [] + + for proc in kill_list: + try: + proc.kill() + except psutil.NoSuchProcess: + pass + + def wait_for_all() -> None: + for proc in kill_list: + try: + proc.wait() + except psutil.NoSuchProcess: + pass + + # Wait for process completion on the executor to avoid blocking the + # event loop. + loop = asyncio.get_event_loop() + wait_for_children = loop.run_in_executor(None, wait_for_all) + + # Send a kill signal to the main process before waiting for the children + # to complete. + try: + process.kill() + await process.wait() + except ProcessLookupError: + _LOG.debug( + 'Process completed before it could be killed. ' + 'This may have been caused by the killing one of its ' + 'child subprocesses.', + ) + + await wait_for_children + + async def run_async( program: str, *args: str, @@ -93,6 +151,20 @@ async def run_async( ) -> CompletedProcess: """Runs a command, capturing and optionally logging its output. + Args: + program: The program to run in a new process. + args: The arguments to pass to the program. + env: An optional mapping of environment variables within which to run + the process. + log_output: Whether to log stdout and stderr of the process to this + process's stdout (prefixed with the PID of the subprocess from which + the output originated). If unspecified, the child process's stdout + and stderr will be captured, and both will be stored in the returned + `CompletedProcess`'s output`. + timeout: An optional floating point number of seconds to allow the + subprocess to run before killing it and its children. If unspecified, + the subprocess will be allowed to continue exiting until it completes. + Returns a CompletedProcess with details from the process. """ @@ -123,13 +195,12 @@ async def run_async( await asyncio.wait_for(process.wait(), timeout) except asyncio.TimeoutError: _LOG.error('%s timed out after %d seconds', program, timeout) - process.kill() - await process.wait() + await _kill_process_and_children(process) if process.returncode: _LOG.error('%s exited with status %d', program, process.returncode) else: - _LOG.debug('%s exited successfully', program) + _LOG.error('%s exited successfully', program) return CompletedProcess(process, output) diff --git a/pw_cli/py/setup.cfg b/pw_cli/py/setup.cfg index 17eeaf9f21..f5f12f9acd 100644 --- a/pw_cli/py/setup.cfg +++ b/pw_cli/py/setup.cfg @@ -19,6 +19,7 @@ author_email = pigweed-developers@googlegroups.com description = Pigweed swiss-army knife [options] +install_requires = psutil packages = find: zip_safe = False install_requires = diff --git a/pw_env_setup/py/pw_env_setup/virtualenv_setup/constraint.list b/pw_env_setup/py/pw_env_setup/virtualenv_setup/constraint.list index 3d11a4cc30..c074f511ba 100644 --- a/pw_env_setup/py/pw_env_setup/virtualenv_setup/constraint.list +++ b/pw_env_setup/py/pw_env_setup/virtualenv_setup/constraint.list @@ -46,6 +46,7 @@ platformdirs==3.0.0 pickleshare==0.7.5 prompt-toolkit==3.0.36 protobuf==3.20.1 +psutil==5.9.4 ptpython==3.0.22 ptyprocess==0.7.0 pyasn1==0.4.8 diff --git a/pw_env_setup/py/pw_env_setup/virtualenv_setup/pigweed_upstream_requirements.txt b/pw_env_setup/py/pw_env_setup/virtualenv_setup/pigweed_upstream_requirements.txt index 72a71b03ca..447f5f1664 100644 --- a/pw_env_setup/py/pw_env_setup/virtualenv_setup/pigweed_upstream_requirements.txt +++ b/pw_env_setup/py/pw_env_setup/virtualenv_setup/pigweed_upstream_requirements.txt @@ -20,5 +20,5 @@ sphinx-design==0.3.0 myst-parser==0.18.1 breathe==4.34.0 # Renode requirements -psutil==5.9.1 +psutil==5.9.4 robotframework==5.0.1