Gracefully shutdown child processes

nox/popen.py

@@ -12,10 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import locale
 import subprocess
 import sys
-from typing import IO, Mapping, Sequence, Tuple, Union
+from typing import IO, Mapping, Optional, Sequence, Tuple, Union
+
+
+def shutdown_process(proc: subprocess.Popen) -> Tuple[Optional[bytes], Optional[bytes]]:
+    """Gracefully shutdown a child process."""
+
+    with contextlib.suppress(subprocess.TimeoutExpired):
+        return proc.communicate(timeout=0.3)
+
+    proc.terminate()
+
+    with contextlib.suppress(subprocess.TimeoutExpired):
+        return proc.communicate(timeout=0.2)
+
+    proc.kill()
+
+    return proc.communicate()
 
 
 def decode_output(output: bytes) -> str:
@@ -57,9 +74,9 @@ def popen(
         sys.stdout.flush()
 
     except KeyboardInterrupt:
-        proc.terminate()
-        proc.wait()
-        raise
+        out, err = shutdown_process(proc)
+        if proc.returncode != 0:
+            raise
 
     return_code = proc.wait()
 

tests/test_command.py

@@ -12,9 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import ctypes
 import logging
 import os
+import platform
+import signal
+import subprocess
 import sys
+import time
+from textwrap import dedent
 from unittest import mock
 
 import nox.command
@@ -23,6 +29,15 @@
 
 PYTHON = sys.executable
 
+skip_on_windows_console = pytest.mark.skipif(
+    platform.system() == "Windows" and "DETACHED_FROM_CONSOLE" not in os.environ,
+    reason="Do not run this test attached to a Windows console.",
+)
+
+only_on_windows = pytest.mark.skipif(
+    platform.system() != "Windows", reason="Only run this test on Windows."
+)
+
 
 def test_run_defaults(capsys):
     result = nox.command.run([PYTHON, "-c", "print(123)"])
@@ -98,7 +113,7 @@ def test_run_env_unicode():
     result = nox.command.run(
         [PYTHON, "-c", 'import os; print(os.environ["SIGIL"])'],
         silent=True,
-        env={u"SIGIL": u"123"},
+        env={"SIGIL": "123"},
     )
 
     assert "123" in result
@@ -196,13 +211,153 @@ def test_fail_with_silent(capsys):
         assert "err" in err
 
 
-def test_interrupt():
-    mock_proc = mock.Mock()
-    mock_proc.communicate.side_effect = KeyboardInterrupt()
+@pytest.fixture
+def marker(tmp_path):
+    """A marker file for process communication."""
+    return tmp_path / "marker"
+
+
+def enable_ctrl_c(enabled):
+    """Enable keyboard interrupts (CTRL-C) on Windows."""
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+
+    if not kernel32.SetConsoleCtrlHandler(None, not enabled):
+        raise ctypes.WinError(ctypes.get_last_error())
+
+
+def interrupt_process(proc):
+    """Send SIGINT or CTRL_C_EVENT to the process."""
+    if platform.system() == "Windows":
+        # Disable Ctrl-C so we don't terminate ourselves.
+        enable_ctrl_c(False)
+
+        # Send the keyboard interrupt to all processes attached to the current
+        # console session.
+        os.kill(0, signal.CTRL_C_EVENT)
+    else:
+        proc.send_signal(signal.SIGINT)
+
+
+@pytest.fixture
+def command_with_keyboard_interrupt(monkeypatch, marker):
+    """Monkeypatch Popen.communicate to raise KeyboardInterrupt."""
+    if platform.system() == "Windows":
+        # Enable Ctrl-C because the child inherits the setting from us.
+        enable_ctrl_c(True)
+
+    communicate = subprocess.Popen.communicate
+
+    def wrapper(proc, *args, **kwargs):
+        # Raise the interrupt only on the first call, so Nox has a chance to
+        # shut down the child process subsequently.
+
+        if wrapper.firstcall:
+            wrapper.firstcall = False
+
+            # Give the child time to install its signal handlers.
+            while not marker.exists():
+                time.sleep(0.05)
+
+            # Send a real keyboard interrupt to the child.
+            interrupt_process(proc)
+
+            # Fake a keyboard interrupt in the parent.
+            raise KeyboardInterrupt
+
+        return communicate(proc, *args, **kwargs)
+
+    wrapper.firstcall = True
+
+    monkeypatch.setattr("subprocess.Popen.communicate", wrapper)
+
+
+def format_program(program, marker):
+    """Preprocess the Python program run by the child process."""
+    main = f"""
+    import time
+    from pathlib import Path
+
+    Path({str(marker)!r}).touch()
+    time.sleep(3)
+    """
+    return dedent(program).format(MAIN=dedent(main))
+
+
+def run_pytest_without_console_window(test):
+    """Run the given test in a subprocess detached from the console window."""
+    env = dict(os.environ, DETACHED_FROM_CONSOLE="")
+    creationflags = (
+        subprocess.CREATE_NO_WINDOW
+        if sys.version_info[:2] >= (3, 7)
+        else subprocess.CREATE_NEW_CONSOLE
+    )
+
+    subprocess.run(
+        ["pytest", f"{__file__}::{test}"],
+        env=env,
+        check=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        creationflags=creationflags,
+    )
+
+
+@skip_on_windows_console
+@pytest.mark.parametrize(
+    "program",
+    [
+        """
+        {MAIN}
+        """,
+        """
+        import signal
+
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+        {MAIN}
+        """,
+        """
+        import signal
+
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+
+        {MAIN}
+        """,
+    ],
+)
+def test_interrupt_raises(command_with_keyboard_interrupt, program, marker):
+    """It kills the process and reraises the keyboard interrupt."""
+    with pytest.raises(KeyboardInterrupt):
+        nox.command.run([PYTHON, "-c", format_program(program, marker)])
+
+
+@only_on_windows
+def test_interrupt_raises_on_windows():
+    """It kills the process and reraises the keyboard interrupt."""
+    run_pytest_without_console_window("test_interrupt_raises")
+
+
+@skip_on_windows_console
+def test_interrupt_handled(command_with_keyboard_interrupt, marker):
+    """It does not raise if the child handles the keyboard interrupt."""
+    program = """
+    import signal
+
+    def exithandler(sig, frame):
+        raise SystemExit()
+
+    signal.signal(signal.SIGINT, exithandler)
+
+    {MAIN}
+    """
+    nox.command.run([PYTHON, "-c", format_program(program, marker)])
+
 
-    with mock.patch("subprocess.Popen", return_value=mock_proc):
-        with pytest.raises(KeyboardInterrupt):
-            nox.command.run([PYTHON, "-c" "123"])
+@only_on_windows
+def test_interrupt_handled_on_windows():
+    """It does not raise if the child handles the keyboard interrupt."""
+    run_pytest_without_console_window("test_interrupt_handled")
 
 
 def test_custom_stdout(capsys, tmpdir):