Merge pull request #49 from ceph/ctrl-c

exec: Clean up subprocesses on ctrl-C

Author: zmc

ceph_devstack/cli.py

@@ -48,9 +48,11 @@ def main() -> int:
     obj = CephDevStack()
 
     async def run():
-        if not await asyncio.gather(
-            check_requirements(),
-            obj.check_requirements(),
+        if not all(
+            await asyncio.gather(
+                check_requirements(),
+                obj.check_requirements(),
+            )
         ):
             logger.error("Requirements not met!")
             return 1

ceph_devstack/exec.py

@@ -1,26 +1,74 @@
 import asyncio
 from asyncio.subprocess import SubprocessStreamProtocol
+import contextlib
 import functools
 import os
 import pathlib
+import psutil
+import signal
 import subprocess
 
 from typing import Dict, List, Optional
 
 from ceph_devstack import logger, VERBOSE
 
+_TERMINATE_TIMEOUT = 3.0
+_KILL_TIMEOUT = 1.0
 
-class LoggingStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
-    def __init__(self, limit, loop, log_level):
-        self.log_level = log_level
-        super().__init__(limit=limit, loop=loop)
 
-    def pipe_data_received(self, fd, data):
-        logger.log(
-            self.log_level,
-            (data.decode() if isinstance(data, bytes) else str(data)).rstrip("\n"),
-        )
-        super().pipe_data_received(fd, data)
+class Subprocess(asyncio.subprocess.Process):
+    async def _close_transport(self) -> None:
+        transport = getattr(self, "_transport", None)
+        if transport is not None and not transport.is_closing():
+            transport.close()
+
+    async def _wait_for_exit(self, timeout: float) -> None:
+        with contextlib.suppress(asyncio.TimeoutError, asyncio.CancelledError):
+            await asyncio.wait_for(asyncio.shield(super().wait()), timeout=timeout)
+
+    async def _terminate(self) -> None:
+        if self.returncode is not None:
+            await self._close_transport()
+            return
+        self.signal_children(signal.SIGTERM, recursive=True)
+        with contextlib.suppress(ProcessLookupError):
+            self.kill()
+        await self._wait_for_exit(_TERMINATE_TIMEOUT)
+        if self.returncode is None:
+            self.signal_children(signal.SIGKILL, recursive=True)
+        with contextlib.suppress(ProcessLookupError):
+            self.kill()
+        await self._wait_for_exit(_KILL_TIMEOUT)
+        await self._close_transport()
+
+    async def wait(self) -> int:
+        try:
+            return await super().wait()
+        except asyncio.CancelledError:
+            await self._terminate()
+            raise
+
+    async def communicate(self, input=None):
+        try:
+            return await super().communicate(input)
+        except asyncio.CancelledError:
+            await self._terminate()
+            raise
+
+    def child_pids(self, recursive=True):
+        if self.pid is None:
+            return []
+        return [
+            child.pid
+            for child in psutil.Process(self.pid).children(recursive=recursive)
+        ]
+
+    def signal_children(self, signal: signal.Signals, recursive=True):
+        for pid in self.child_pids(recursive=recursive):
+            with contextlib.suppress(ProcessLookupError):
+                os.kill(pid, signal)
+            with contextlib.suppress(ProcessLookupError):
+                os.killpg(pid, signal)
 
 
 class Command:
@@ -57,37 +105,28 @@ def run(self) -> subprocess.Popen:
         proc.wait()
         return proc
 
-    async def arun(self) -> asyncio.subprocess.Process:
+    async def arun(self) -> Subprocess:
         logger.log(VERBOSE, self._make_log_msg())
         loop = asyncio.get_running_loop()
-        protocol_factory: (
-            functools.partial[SubprocessStreamProtocol]
-            | functools.partial[LoggingStreamProtocol]
-        )
+        kwargs = dict(self.kwargs)
         if self.stream_output:
-            protocol_factory = functools.partial(
-                LoggingStreamProtocol,
-                limit=2**16,
-                loop=loop,
-                log_level=VERBOSE,
-            )
-        else:
-            protocol_factory = functools.partial(
-                asyncio.subprocess.SubprocessStreamProtocol,
-                limit=2**16,
-                loop=loop,
-            )
+            # Inherit stdout/stderr so long-running commands (e.g. dnf builddep)
+            # are not blocked when their output exceeds the StreamReader limit.
+            kwargs["stdout"] = None
+            kwargs["stderr"] = None
+        protocol_factory = functools.partial(
+            SubprocessStreamProtocol,
+            limit=2**16,
+            loop=loop,
+        )
         transport, protocol = await loop.subprocess_exec(
             protocol_factory,
             *self.args,
             env=self.env,
-            **self.kwargs,
-        )
-        return asyncio.subprocess.Process(
-            transport,
-            protocol,
-            loop,
+            start_new_session=True,
+            **kwargs,
         )
+        return Subprocess(transport, protocol, loop)
 
     def __str__(self):
         return " ".join(self.args)

ceph_devstack/host.py

@@ -1,4 +1,3 @@
-import asyncio
 import json
 import logging
 import os
@@ -10,7 +9,7 @@
 from packaging.version import parse as parse_version, Version
 from typing import Dict, List, Optional, Union
 
-from .exec import Command
+from .exec import Command, Subprocess
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +45,7 @@ async def arun(
         cwd: Optional[pathlib.Path] = None,
         env: Optional[Dict] = None,
         stream_output: bool = False,
-    ) -> asyncio.subprocess.Process:
+    ) -> Subprocess:
         return await self.cmd(
             args, cwd=cwd, env=env, stream_output=stream_output
         ).arun()
@@ -145,7 +144,15 @@ class LocalHost(Host):
 
 class RemoteHost(Host):
     type = "remote"
-    base_args = ["podman", "machine", "ssh", "--"]
+    base_args = ["podman", "machine", "ssh"]
+
+    def _remote_args(self, args: List[str], stream_output: bool) -> List[str]:
+        remote = list(self.base_args)
+        if stream_output:
+            remote.append("-t")
+        remote.append("--")
+        remote.extend(args)
+        return remote
 
     def cmd(
         self,
@@ -155,7 +162,7 @@ def cmd(
         stream_output: bool = False,
     ):
         if args[0] != "podman":
-            args = self.base_args + args
+            args = self._remote_args(args, stream_output)
         return super().cmd(args, cwd=cwd, env=env, stream_output=stream_output)
 
     def path_exists(self, path: Union[str, pathlib.Path]):

ceph_devstack/resources/__init__.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import argparse
-import asyncio
 import json
 import os
 import subprocess
@@ -9,6 +8,7 @@
 from subprocess import CalledProcessError
 from typing import List, Dict, Set
 
+from ceph_devstack.exec import Subprocess
 from ceph_devstack.host import host, local_host
 
 
@@ -56,15 +56,13 @@ async def cmd(
         check: bool = False,
         force_local: bool = False,
         stream_output: bool = False,
-    ) -> asyncio.subprocess.Process:
+    ) -> Subprocess:
         exec_host = local_host if force_local else host
         proc = await exec_host.arun(
             args,
             cwd=Path(self.cwd),
             stream_output=stream_output,
         )
-        assert proc.stderr is not None
-        assert proc.stdout is not None
         returncode = await proc.wait()
         if check and returncode != 0:
             # out = await proc.stderr.read()

ceph_devstack/resources/ceph/__init__.py

@@ -205,22 +205,17 @@ async def watch(self):
                 containers.append(object)
         logger.info(f"Watching {containers}")
         while True:
-            try:
-                for container in containers:
-                    with contextlib.suppress(CalledProcessError):
-                        if not await container.exists():
-                            logger.info(
-                                f"Container {container.name} was removed; replacing"
-                            )
-                            await container.create()
-                            await container.start()
-                        elif not await container.is_running():
-                            logger.info(
-                                f"Container {container.name} stopped; restarting"
-                            )
-                            await container.start()
-            except KeyboardInterrupt:
-                break
+            for container in containers:
+                with contextlib.suppress(CalledProcessError):
+                    if not await container.exists():
+                        logger.info(
+                            f"Container {container.name} was removed; replacing"
+                        )
+                        await container.create()
+                        await container.start()
+                    elif not await container.is_running():
+                        logger.info(f"Container {container.name} stopped; restarting")
+                        await container.start()
 
     async def wait(self, container_name: str):
         for spec in self.service_specs.values():

pyproject.toml

@@ -13,7 +13,13 @@ classifiers = [
 keywords = ["podman", "ceph"]
 description = "Run a full teuthology lab on your laptop!"
 requires-python = ">=3.12"
-dependencies = ["packaging", "pre-commit", "PyYAML", "tomlkit"]
+dependencies = [
+    "packaging",
+    "pre-commit",
+    "psutil>=7.2.2",
+    "PyYAML",
+    "tomlkit",
+]
 dynamic = ["version"]
 
 [project.readme]