nixos-rebuild-ng: cleanup SSH ControlMaster at exit

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

@@ -1,4 +1,5 @@
 import argparse
+import atexit
 import json
 import os
 import sys
@@ -22,7 +23,7 @@ from .nix import (
     switch_to_configuration,
     upgrade_channels,
 )
-from .process import Remote
+from .process import Remote, cleanup_ssh
 from .utils import flags_to_dict, info
 
 VERBOSE = 0
@@ -176,6 +177,7 @@ def execute(argv: list[str]) -> None:
     # Will be cleaned up on exit automatically.
     tmpdir = TemporaryDirectory(prefix="nixos-rebuild.")
     tmpdir_path = Path(tmpdir.name)
+    atexit.register(cleanup_ssh, tmpdir_path)
 
     profile = Profile.from_name(args.profile_name)
     target_host = Remote.from_arg(args.target_host, not args.no_ssh_tty, tmpdir_path)

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

@@ -26,6 +26,7 @@ class Remote:
     def from_arg(cls, host: str | None, tty: bool | None, tmp_dir: Path) -> Self | None:
         if host:
             opts = os.getenv("NIX_SSHOPTS", "").split() + [
+                # SSH ControlMaster flags, allow for faster re-connection
                 "-o",
                 "ControlMaster=auto",
                 "-o",
@@ -38,6 +39,12 @@ class Remote:
             return None
 
 
+def cleanup_ssh(tmp_dir: Path) -> None:
+    "Close SSH ControlMaster connection."
+    for ctrl in tmp_dir.glob("ssh-*"):
+        subprocess.run(["ssh", "-o", f"ControlPath={ctrl}", "exit"], check=False)
+
+
 def run_wrapper(
     args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]],
     *,