From 60ecb14418d49b2f5d16bcc6729b376908df24bb Mon Sep 17 00:00:00 2001 From: A5rocks Date: Tue, 27 May 2025 13:28:56 +0900 Subject: [PATCH 1/2] Implement helpful features from https://github.com/python-trio/trio/issues/3237 - `shell=True` implies `process_group=0` - make default `deliver_cancel` aware of process groups - backport `process_group` to Pythons prior to 3.11 --- newsfragments/3237.feature.rst | 1 + src/trio/_subprocess.py | 88 +++++++++++++++++++++++------- src/trio/_tests/test_subprocess.py | 82 ++++++++++++++++++++++++++++ 3 files changed, 152 insertions(+), 19 deletions(-) create mode 100644 newsfragments/3237.feature.rst diff --git a/newsfragments/3237.feature.rst b/newsfragments/3237.feature.rst new file mode 100644 index 0000000000..d74952c641 --- /dev/null +++ b/newsfragments/3237.feature.rst @@ -0,0 +1 @@ +Improve defaults if ``shell=True`` for `trio.run_process`, as well as backport the ``process_group`` argument. diff --git a/src/trio/_subprocess.py b/src/trio/_subprocess.py index c1b162e308..02c5dd6f14 100644 --- a/src/trio/_subprocess.py +++ b/src/trio/_subprocess.py @@ -2,9 +2,11 @@ import contextlib import os +import signal import subprocess import sys import warnings +from collections.abc import Awaitable from contextlib import ExitStack from functools import partial from typing import ( @@ -30,7 +32,6 @@ from ._util import NoPublicConstructor, final if TYPE_CHECKING: - import signal from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from io import TextIOWrapper @@ -441,24 +442,45 @@ async def _windows_deliver_cancel(p: Process) -> None: # noqa: RUF029 ) -async def _posix_deliver_cancel(p: Process) -> None: - try: - p.terminate() - await trio.sleep(5) - warnings.warn( - RuntimeWarning( - f"process {p!r} ignored SIGTERM for 5 seconds. " - "(Maybe you should pass a custom deliver_cancel?) " - "Trying SIGKILL.", - ), - stacklevel=1, - ) - p.kill() - except OSError as exc: - warnings.warn( - RuntimeWarning(f"tried to kill process {p!r}, but failed with: {exc!r}"), - stacklevel=1, +def _get_posix_deliver_cancel( + process_group: int | None, +) -> Callable[[Process], Awaitable[None]]: + async def _posix_deliver_cancel(p: Process) -> None: + should_deliver_to_pg = ( + sys.platform != "win32" + and process_group is not None + and os.getpgrp() != os.getpgid(p.pid) ) + try: + # TODO: should Process#terminate do this special logic + if sys.platform != "win32" and should_deliver_to_pg: + os.killpg(os.getpgid(p.pid), signal.SIGTERM) + else: + p.terminate() + + await trio.sleep(5) + warnings.warn( + RuntimeWarning( + f"process {p!r} ignored SIGTERM for 5 seconds. " + "(Maybe you should pass a custom deliver_cancel?) " + "Trying SIGKILL.", + ), + stacklevel=1, + ) + + if sys.platform != "win32" and should_deliver_to_pg: + os.killpg(os.getpgid(p.pid), signal.SIGKILL) + else: + p.kill() + except OSError as exc: + warnings.warn( + RuntimeWarning( + f"tried to kill process {p!r}, but failed with: {exc!r}" + ), + stacklevel=1, + ) + + return _posix_deliver_cancel # Use a private name, so we can declare platform-specific stubs below. @@ -472,6 +494,7 @@ async def _run_process( check: bool = True, deliver_cancel: Callable[[Process], Awaitable[object]] | None = None, task_status: TaskStatus[Process] = trio.TASK_STATUS_IGNORED, + shell: bool = False, **options: object, ) -> subprocess.CompletedProcess[bytes]: """Run ``command`` in a subprocess and wait for it to complete. @@ -689,6 +712,9 @@ async def my_deliver_cancel(process): "stderr=subprocess.PIPE is only valid with nursery.start, " "since that's the only way to access the pipe", ) + + options["shell"] = shell + if isinstance(stdin, (bytes, bytearray, memoryview)): input_ = stdin options["stdin"] = subprocess.PIPE @@ -708,12 +734,36 @@ async def my_deliver_cancel(process): raise ValueError("can't specify both stderr and capture_stderr") options["stderr"] = subprocess.PIPE + # ensure things can be killed including a shell's child processes + if shell and sys.platform != "win32" and "process_group" not in options: + options["process_group"] = 0 + if deliver_cancel is None: if os.name == "nt": deliver_cancel = _windows_deliver_cancel else: assert os.name == "posix" - deliver_cancel = _posix_deliver_cancel + deliver_cancel = _get_posix_deliver_cancel(options.get("process_group")) # type: ignore[arg-type] + + if ( + sys.platform != "win32" + and options.get("process_group") is not None + and sys.version_info < (3, 11) + ): + # backport the argument to Python versions prior to 3.11 + preexec_fn = options.get("preexec_fn") + process_group = options.pop("process_group") + + def new_preexecfn() -> object: # pragma: no cover + assert sys.platform != "win32" + os.setpgid(0, process_group) # type: ignore[arg-type] + + if callable(preexec_fn): + return preexec_fn() + else: + return None + + options["preexec_fn"] = new_preexecfn stdout_chunks: list[bytes | bytearray] = [] stderr_chunks: list[bytes | bytearray] = [] diff --git a/src/trio/_tests/test_subprocess.py b/src/trio/_tests/test_subprocess.py index 15c88be25e..1f1e22edce 100644 --- a/src/trio/_tests/test_subprocess.py +++ b/src/trio/_tests/test_subprocess.py @@ -765,3 +765,85 @@ async def wait_and_tell(proc: Process) -> None: # for everything to notice await noticed_exit.wait() assert noticed_exit.is_set(), "child task wasn't woken after poll, DEADLOCK" + + +@pytest.mark.skipif(not posix, reason="posix only") +@slow +async def test_shells_killed_by_default() -> None: + assert sys.platform != "win32" + + async with _core.open_nursery() as nursery: + proc = await nursery.start( + partial( + trio.run_process, + 'python -u -c "import os, time, signal; ' + f"os.kill({os.getpid()}, signal.SIGHUP); " + 'print(os.getpid()); time.sleep(5)" | cat', + shell=True, + stdout=subprocess.PIPE, + ) + ) + + with trio.fail_after(1): + with trio.open_signal_receiver(signal.SIGHUP) as signal_aiter: # type: ignore[attr-defined,unused-ignore] + async for signum in signal_aiter: # pragma: no branch + assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] + break + + nursery.cancel_scope.cancel() + + chunks = [] + with trio.move_on_after(0.1): + async with proc.stdout: + async for c in proc.stdout: + chunks.append(c) # noqa: PERF401 + + # give the child time to wind down execution (this is unfortunate) + await trio.sleep(1) + + child_pid = int(b"".join(chunks)) + if sys.platform == "linux": + # docker containers can have an evil init which + # won't reap zombie children, so we have to work around that + with pytest.raises(FileNotFoundError): + with open(f"/proc/{child_pid}/status") as f: # noqa: ASYNC230 + for line in f: # pragma: no branch + if line.startswith("State:"): + assert "Z" in line + raise FileNotFoundError + else: + with pytest.raises(OSError, match="No such process"): + os.kill(child_pid, 0) + + +@pytest.mark.skipif(not posix, reason="posix only") +@slow +async def test_process_group_SIGKILL_escalation(mock_clock: MockClock) -> None: + assert sys.platform != "win32" + mock_clock.autojump_threshold = 1 + + with pytest.warns(RuntimeWarning, match=".*ignored SIGTERM.*"): # noqa: PT031 + async with _core.open_nursery() as nursery: + nursery.start_soon( + # this ignore is because process_group=0 on Windows + partial( # type: ignore[call-arg,unused-ignore] + trio.run_process, + [ + "python", + "-c", + "import os, time, signal;" + f"os.kill({os.getpid()}, signal.SIGHUP);" + "signal.signal(signal.SIGTERM, signal.SIG_IGN);" + "time.sleep(10)", + ], + process_group=0, + ) + ) + + with trio.fail_after(1): + with trio.open_signal_receiver(signal.SIGHUP) as signal_aiter: # type: ignore[attr-defined,unused-ignore] + async for signum in signal_aiter: # pragma: no branch + assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] + break + + nursery.cancel_scope.cancel() From ee0cd2554fe26edc9882adbc49d4a4531dc007e7 Mon Sep 17 00:00:00 2001 From: A5rocks Date: Tue, 27 May 2025 13:42:34 +0900 Subject: [PATCH 2/2] Clean up coverage pragmas --- src/trio/_tests/test_subprocess.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/trio/_tests/test_subprocess.py b/src/trio/_tests/test_subprocess.py index 1f1e22edce..7f266a927e 100644 --- a/src/trio/_tests/test_subprocess.py +++ b/src/trio/_tests/test_subprocess.py @@ -786,9 +786,8 @@ async def test_shells_killed_by_default() -> None: with trio.fail_after(1): with trio.open_signal_receiver(signal.SIGHUP) as signal_aiter: # type: ignore[attr-defined,unused-ignore] - async for signum in signal_aiter: # pragma: no branch - assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] - break + signum = await signal_aiter.__anext__() + assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] nursery.cancel_scope.cancel() @@ -807,10 +806,14 @@ async def test_shells_killed_by_default() -> None: # won't reap zombie children, so we have to work around that with pytest.raises(FileNotFoundError): with open(f"/proc/{child_pid}/status") as f: # noqa: ASYNC230 - for line in f: # pragma: no branch + hit_zombie = False + for line in f: if line.startswith("State:"): assert "Z" in line - raise FileNotFoundError + hit_zombie = True + + if hit_zombie: + raise FileNotFoundError else: with pytest.raises(OSError, match="No such process"): os.kill(child_pid, 0) @@ -842,8 +845,7 @@ async def test_process_group_SIGKILL_escalation(mock_clock: MockClock) -> None: with trio.fail_after(1): with trio.open_signal_receiver(signal.SIGHUP) as signal_aiter: # type: ignore[attr-defined,unused-ignore] - async for signum in signal_aiter: # pragma: no branch - assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] - break + signum = await signal_aiter.__anext__() + assert signum == signal.SIGHUP # type: ignore[attr-defined,unused-ignore] nursery.cancel_scope.cancel()