From 927d182d908eee6648e4e0b237a78f3f31cd697d Mon Sep 17 00:00:00 2001
From: Thiago Kenji Okada <thiagokokada@gmail.com>
Date: Mon, 16 Dec 2024 11:12:38 +0000
Subject: [PATCH] nixos-rebuild-ng: move TMPDIR to its own file

---
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py |  5 ++++-
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/process.py  | 11 ++++-------
 .../ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py   |  5 +++++
 3 files changed, 13 insertions(+), 8 deletions(-)
 create mode 100644 pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py

diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index 50c4d9ba787f..ab406e7e0d87 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -8,7 +8,7 @@ from pathlib import Path
 from subprocess import CalledProcessError, run
 from typing import assert_never
 
-from . import nix
+from . import nix, tmpdir
 from .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES
 from .models import Action, BuildAttr, Flake, NRError, Profile
 from .process import Remote, cleanup_ssh
@@ -280,7 +280,10 @@ def reexec(
                 argv[0],
                 new,
             )
+            # Manually call clean-up functions since os.execve() will replace
+            # the process immediately
             cleanup_ssh()
+            tmpdir.TMPDIR.cleanup()
             os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
 
 
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
index 10666b47d657..337fae18b2ff 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
@@ -4,19 +4,17 @@ import shlex
 import subprocess
 from dataclasses import dataclass
 from getpass import getpass
-from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Self, Sequence, TypedDict, Unpack
 
+from . import tmpdir
+
 logger = logging.getLogger(__name__)
 
-TMPDIR = TemporaryDirectory(prefix="nixos-rebuild.")
-TMPDIR_PATH = Path(TMPDIR.name)
 SSH_DEFAULT_OPTS = [
     "-o",
     "ControlMaster=auto",
     "-o",
-    f"ControlPath={TMPDIR_PATH / "ssh-%n"}",
+    f"ControlPath={tmpdir.TMPDIR_PATH / "ssh-%n"}",
     "-o",
     "ControlPersist=60",
 ]
@@ -70,13 +68,12 @@ class RunKwargs(TypedDict, total=False):
 
 def cleanup_ssh() -> None:
     "Close SSH ControlMaster connection."
-    for ctrl in TMPDIR_PATH.glob("ssh-*"):
+    for ctrl in tmpdir.TMPDIR_PATH.glob("ssh-*"):
         run_wrapper(
             ["ssh", "-o", f"ControlPath={ctrl}", "-O", "exit", "dummyhost"],
             check=False,
             capture_output=True,
         )
-    TMPDIR.cleanup()
 
 
 def run_wrapper(
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py
new file mode 100644
index 000000000000..ca71a6ddd8c3
--- /dev/null
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/tmpdir.py
@@ -0,0 +1,5 @@
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+TMPDIR = TemporaryDirectory(prefix="nixos-rebuild.")
+TMPDIR_PATH = Path(TMPDIR.name)