nixos-rebuild-ng: kill underlying remote process

`nixos-rebuild-ng` explicitly don't allocate a pseudo-TTY for SSH
because this causes lots of issues depending on the use case (for
example, multiplexing multiple SSH sessions).

Sadly not using a pseudo-TTY also cause other issues, like the fact that
using Ctrl+C (SIGINT) doesn't kill the underlying process because SSH
doesn't support it.

We can't really start using pseudo-TTY unless we want to overcomplicate
the code for parsing results (pseudo-TTY mangles the stdout/stderr
together), so we need to handle killing the underlying remote process
manually.

This is what this commit does, when we receive a `KeyboardInterrupt`
exception while calling `run_wrapper`, we will check if it is a remote
process and send a `pkill --full` with the arguments (this should ensure
that we don't kill other process, but we can't guarantee it). This
assumes the user has `procps` installed, but I think it is a safe
assumption since this seems to be a core package.

Sadly nothing we can do if the user doesn't have `procps` installed, the
good thing is that the worst that can happen is that we will silent
fail and the process will stay in background until it finishes.

Fix #403269.

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py

@@ -1,6 +1,7 @@
 import atexit
 import logging
 import os
+import re
 import shlex
 import subprocess
 from collections.abc import Sequence
@@ -21,6 +22,8 @@ SSH_DEFAULT_OPTS: Final = [
     "ControlPersist=60",
 ]
 
+type Args = Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]]
+
 
 @dataclass(frozen=True)
 class Remote:
@@ -82,7 +85,7 @@ atexit.register(cleanup_ssh)
 
 
 def run_wrapper(
-    args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]],
+    args: Args,
     *,
     check: bool = True,
     extra_env: dict[str, str] | None = None,
@@ -93,6 +96,8 @@ def run_wrapper(
     "Wrapper around `subprocess.run` that supports extra functionality."
     env = None
     process_input = None
+    run_args = args
+
     if remote:
         if extra_env:
             extra_env_args = [f"{env}={value}" for env, value in extra_env.items()]
@@ -103,7 +108,7 @@ def run_wrapper(
                 process_input = remote.sudo_password + "\n"
             else:
                 args = ["sudo", *args]
-        args = [
+        run_args = [
             "ssh",
             *remote.opts,
             *SSH_DEFAULT_OPTS,
@@ -119,32 +124,39 @@ def run_wrapper(
         if extra_env:
             env = os.environ | extra_env
         if sudo:
-            args = ["sudo", *args]
+            run_args = ["sudo", *run_args]
 
     logger.debug(
         "calling run with args=%r, kwargs=%r, extra_env=%r",
-        args,
+        run_args,
         kwargs,
         extra_env,
     )
 
     try:
         r = subprocess.run(
-            args,
+            run_args,
             check=check,
             env=env,
             input=process_input,
-            # Hope nobody is using NixOS with non-UTF8 encodings, but "surrogateescape"
-            # should still work in those systems.
+            # Hope nobody is using NixOS with non-UTF8 encodings, but
+            # "surrogateescape" should still work in those systems.
             text=True,
             errors="surrogateescape",
             **kwargs,
         )
 
         if kwargs.get("capture_output") or kwargs.get("stderr") or kwargs.get("stdout"):
-            logger.debug("captured output stdout=%r, stderr=%r", r.stdout, r.stderr)
+            logger.debug(
+                "captured output with stdout=%r, stderr=%r", r.stdout, r.stderr
+            )
 
         return r
+    except KeyboardInterrupt:
+        # sudo commands are activation only and unlikely to be long running
+        if remote and not sudo:
+            _kill_long_running_ssh_process(args, remote)
+        raise
     except subprocess.CalledProcessError:
         if sudo and remote and remote.sudo_password is None:
             logger.error(
@@ -152,3 +164,42 @@ def run_wrapper(
                 + "--ask-sudo-password?"
             )
         raise
+
+
+# SSH does not send the signals to the process when running without usage of
+# pseudo-TTY (that causes a whole other can of worms), so if the process is
+# long running (e.g.: a build) this will result in the underlying process
+# staying alive.
+# See: https://stackoverflow.com/a/44354466
+# Issue: https://github.com/NixOS/nixpkgs/issues/403269
+def _kill_long_running_ssh_process(args: Args, remote: Remote) -> None:
+    logger.info("cleaning-up remote process, please wait...")
+
+    # We need to escape both the shell and regex here (since pkill interprets
+    # its arguments as regex)
+    quoted_args = re.escape(shlex.join(str(a) for a in args))
+    logger.debug("killing remote process using pkill with args=%r", quoted_args)
+    r = subprocess.run(
+        [
+            "ssh",
+            *remote.opts,
+            *SSH_DEFAULT_OPTS,
+            remote.host,
+            "--",
+            "pkill",
+            "--signal",
+            "SIGINT",
+            "--full",
+            "--",
+            quoted_args,
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    logger.debug(
+        "remote pkill captured output with stdout=%r, stderr=%r, returncode=%s",
+        r.stdout,
+        r.stderr,
+        r.returncode,
+    )

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_process.py

@@ -96,6 +96,37 @@ def test_run(mock_run: Any) -> None:
     )
 
 
+@patch(get_qualified_name(p.subprocess.run), autospec=True)
+def test__kill_long_running_ssh_process(mock_run: Any) -> None:
+    p._kill_long_running_ssh_process(
+        [
+            "nix",
+            "--extra-experimental-features",
+            "nix-command flakes",
+            "build",
+            "/nix/store/la0c8nmpr9xfclla0n4f3qq9iwgdrq4g-nixos-system-sankyuu-nixos-25.05.20250424.f771eb4.drv^*",
+        ],
+        m.Remote("user@localhost", opts=[], sudo_password=None),
+    )
+    mock_run.assert_called_with(
+        [
+            "ssh",
+            *p.SSH_DEFAULT_OPTS,
+            "user@localhost",
+            "--",
+            "pkill",
+            "--signal",
+            "SIGINT",
+            "--full",
+            "--",
+            r"nix\ \-\-extra\-experimental\-features\ 'nix\-command\ flakes'\ build\ '/nix/store/la0c8nmpr9xfclla0n4f3qq9iwgdrq4g\-nixos\-system\-sankyuu\-nixos\-25\.05\.20250424\.f771eb4\.drv\^\*'",
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+
 def test_remote_from_name(monkeypatch: MonkeyPatch) -> None:
     monkeypatch.setenv("NIX_SSHOPTS", "")
     assert m.Remote.from_arg("user@localhost", None, False) == m.Remote(