pw_cli: Kill child processes of timed-out run

Tests run using `run`/`run_async` which timed out would previously leave
around their child processes. This could disrupt system state.

After this CL, when a timeout occurs, child processes will be killed and
waited for completion before killing and waiting for the completion of
the main process.

Fixed: b/243591004
Change-Id: Ifeb0796b45c13087cede1b05d2971c195c30607a
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/131527
Commit-Queue: Taylor Cramer <[email protected]>
Reviewed-by: Anthony DiGirolamo <[email protected]>
Reviewed-by: Austin Foxley <[email protected]>

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)",

pw_cli/py/BUILD.bazel

@@ -51,21 +51,21 @@ py_binary(
 )
 
 py_test(
-    name = "plugins_test",
+    name = "envparse_test",
     size = "small",
     srcs = [
-        "plugins_test.py",
+        "envparse_test.py",
     ],
     deps = [
         ":pw_cli",
     ],
 )
 
 py_test(
-    name = "envparse_test",
+    name = "plugins_test",
     size = "small",
     srcs = [
-        "envparse_test.py",
+        "plugins_test.py",
     ],
     deps = [
         ":pw_cli",

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
+  }
+}

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()

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)
 

pw_cli/py/setup.cfg

@@ -19,6 +19,7 @@ author_email = [email protected]
 description = Pigweed swiss-army knife
 
 [options]
+install_requires = psutil
 packages = find:
 zip_safe = False
 install_requires =

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

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