diff --git a/ceph_devstack/cli.py b/ceph_devstack/cli.py index c0794170..8658bcbc 100644 --- a/ceph_devstack/cli.py +++ b/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 diff --git a/ceph_devstack/exec.py b/ceph_devstack/exec.py index bfeb9363..72d4c5c1 100644 --- a/ceph_devstack/exec.py +++ b/ceph_devstack/exec.py @@ -1,26 +1,81 @@ import asyncio from asyncio.subprocess import SubprocessStreamProtocol +import contextlib import functools import os import pathlib +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"), +def _child_pids(pid: int) -> List[int]: + try: + out = subprocess.check_output( + ["pgrep", "-P", str(pid)], + stderr=subprocess.DEVNULL, + text=True, ) - super().pipe_data_received(fd, data) + except (subprocess.CalledProcessError, FileNotFoundError): + return [] + return [int(p) for p in out.split() if p] + + +def _signal_process_tree(pid: int, sig: int) -> None: + for child in _child_pids(pid): + _signal_process_tree(child, sig) + with contextlib.suppress(ProcessLookupError): + os.kill(pid, sig) + with contextlib.suppress(ProcessLookupError): + os.killpg(pid, sig) + + +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 + pid = self.pid + if pid is not None: + _signal_process_tree(pid, signal.SIGTERM) + with contextlib.suppress(ProcessLookupError): + self.kill() + await self._wait_for_exit(_TERMINATE_TIMEOUT) + if self.returncode is None and pid is not None: + _signal_process_tree(pid, signal.SIGKILL) + 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 class Command: @@ -57,37 +112,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) diff --git a/ceph_devstack/host.py b/ceph_devstack/host.py index 5c2f8836..90cf3477 100644 --- a/ceph_devstack/host.py +++ b/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]): diff --git a/ceph_devstack/resources/__init__.py b/ceph_devstack/resources/__init__.py index 65051cc3..0a3c2a10 100644 --- a/ceph_devstack/resources/__init__.py +++ b/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() diff --git a/ceph_devstack/resources/ceph/__init__.py b/ceph_devstack/resources/ceph/__init__.py index 436eb3bb..17cc5641 100644 --- a/ceph_devstack/resources/ceph/__init__.py +++ b/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(): diff --git a/tests/test_exec.py b/tests/test_exec.py new file mode 100644 index 00000000..8874a7fa --- /dev/null +++ b/tests/test_exec.py @@ -0,0 +1,163 @@ +import asyncio +import os +import sys + +import pytest + +from ceph_devstack.exec import Command +from ceph_devstack.host import RemoteHost + + +def test_remote_host_uses_tty_for_streaming(): + host = RemoteHost() + cmd = host.cmd(["python", "build.py"], stream_output=True) + assert cmd.args == ["podman", "machine", "ssh", "-t", "--", "python", "build.py"] + + +def test_remote_host_no_tty_for_buffered_remote_cmd(): + host = RemoteHost() + cmd = host.cmd(["python", "check.py"], stream_output=False) + assert cmd.args == ["podman", "machine", "ssh", "--", "python", "check.py"] + + +def test_remote_host_podman_cmd_not_wrapped(): + host = RemoteHost() + cmd = host.cmd(["podman", "build", "."], stream_output=True) + assert cmd.args == ["podman", "build", "."] + + +@pytest.mark.asyncio +async def test_cancel_kills_streaming_build_process(): + script = "import time; time.sleep(3600)" + cmd = Command([sys.executable, "-c", script], stream_output=True) + proc = await cmd.arun() + pid = proc.pid + task = asyncio.create_task(proc.wait()) + await asyncio.sleep(0.2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + +@pytest.mark.asyncio +async def test_stream_output_writes_stdout_to_terminal(capfd): + cmd = Command( + [sys.executable, "-c", "print('build-output', flush=True)"], + stream_output=True, + ) + proc = await cmd.arun() + assert await proc.wait() == 0 + captured = capfd.readouterr() + assert "build-output" in captured.out + assert captured.err == "" + + +@pytest.mark.asyncio +async def test_stream_output_writes_stderr_to_terminal(capfd): + cmd = Command( + [ + sys.executable, + "-c", + "import sys; print('build-error', file=sys.stderr, flush=True)", + ], + stream_output=True, + ) + proc = await cmd.arun() + assert await proc.wait() == 0 + captured = capfd.readouterr() + assert captured.out == "" + assert "build-error" in captured.err + + +@pytest.mark.asyncio +async def test_stream_output_does_not_deadlock_on_large_output(capfd): + cmd = Command( + [ + sys.executable, + "-c", + "print('x' * 100_000, flush=True); print('done', flush=True)", + ], + stream_output=True, + ) + proc = await cmd.arun() + await asyncio.wait_for(proc.wait(), timeout=5) + captured = capfd.readouterr() + assert "done" in captured.out + + +@pytest.mark.asyncio +async def test_buffered_output_not_written_to_terminal(capfd): + cmd = Command( + [sys.executable, "-c", "print('hidden', flush=True)"], + stream_output=False, + ) + proc = await cmd.arun() + assert await proc.wait() == 0 + captured = capfd.readouterr() + assert captured.out == "" + assert captured.err == "" + assert proc.stdout is not None + assert (await proc.stdout.read()).decode() == "hidden\n" + + +@pytest.mark.asyncio +async def test_subprocess_transport_closed_on_cancel(): + cmd = Command( + [sys.executable, "-c", "import time; time.sleep(3600)"], + stream_output=True, + ) + proc = await cmd.arun() + task = asyncio.create_task(proc.wait()) + await asyncio.sleep(0.2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert proc.returncode is not None + assert proc._transport.is_closing() + + +def test_subprocess_cancel_does_not_error_on_loop_shutdown(): + async def run_build(): + cmd = Command( + [sys.executable, "-c", "import time; time.sleep(3600)"], + stream_output=True, + ) + proc = await cmd.arun() + await proc.wait() + return proc.pid + + with pytest.raises(TimeoutError): + asyncio.run(asyncio.wait_for(run_build(), timeout=0.2)) + + +@pytest.mark.asyncio +async def test_terminate_kills_child_in_process_group(tmp_path): + child_pid_file = tmp_path / "child.pid" + script = ( + "import subprocess, sys, time\n" + f"child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(3600)'])\n" + f"open({str(child_pid_file)!r}, 'w').write(str(child.pid))\n" + "time.sleep(3600)\n" + ) + cmd = Command([sys.executable, "-c", script], stream_output=True) + proc = await cmd.arun() + + child_pid = None + for _ in range(20): + if child_pid_file.is_file(): + child_pid = int(child_pid_file.read_text()) + break + await asyncio.sleep(0.05) + assert child_pid is not None + + task = asyncio.create_task(proc.wait()) + await asyncio.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + with pytest.raises(ProcessLookupError): + os.kill(child_pid, 0)