From b6a08a2fed8a48d3759ea67e958c9f9ec5f44d94 Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Fri, 5 Apr 2024 22:22:44 +0200
Subject: [PATCH] libutil: return a program handle from runProgram2

this will let us also return a source for the program output later,
which will in turn make sinkToSource unnecessary for program output
processing. this may also reopen a path for provigin program input,
but that still needs a proper async io framework to avoid problems.

Change-Id: Iaf93f47db99c38cfaf134bd60ed6a804d7ddf688
---
 src/libcmd/repl.cc                    |  4 ++--
 src/libfetchers/git.cc                |  2 +-
 src/libstore/build/derivation-goal.cc |  2 +-
 src/libutil/processes.cc              | 28 ++++++++++++++++++++-------
 src/libutil/processes.hh              | 19 +++++++++++++++++-
 5 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/src/libcmd/repl.cc b/src/libcmd/repl.cc
index 39e89d999..bb9225ce5 100644
--- a/src/libcmd/repl.cc
+++ b/src/libcmd/repl.cc
@@ -251,7 +251,7 @@ void runNix(Path program, const Strings & args)
         .program = settings.nixBinDir+ "/" + program,
         .args = args,
         .environment = subprocessEnv,
-    });
+    }).wait();
 
     return;
 }
@@ -650,7 +650,7 @@ ProcessLineResult NixRepl::processLine(std::string line)
 
         // runProgram redirects stdout to a StringSink,
         // using runProgram2 to allow editors to display their UI
-        runProgram2(RunOptions { .program = editor, .searchPath = true, .args = args });
+        runProgram2(RunOptions { .program = editor, .searchPath = true, .args = args }).wait();
 
         // Reload right after exiting the editor
         state->resetFileCache();
diff --git a/src/libfetchers/git.cc b/src/libfetchers/git.cc
index 10e125207..75ac926e2 100644
--- a/src/libfetchers/git.cc
+++ b/src/libfetchers/git.cc
@@ -695,7 +695,7 @@ struct GitInputScheme : InputScheme
                     .program = "git",
                     .args = { "-C", repoDir, "--git-dir", gitDir, "archive", input.getRev()->gitRev() },
                     .standardOut = &sink
-                });
+                }).wait();
             });
 
             unpackTarfile(*source, tmpDir);
diff --git a/src/libstore/build/derivation-goal.cc b/src/libstore/build/derivation-goal.cc
index c0cd523e6..a59984688 100644
--- a/src/libstore/build/derivation-goal.cc
+++ b/src/libstore/build/derivation-goal.cc
@@ -928,7 +928,7 @@ void runPostBuildHook(
         .environment = hookEnvironment,
         .standardOut = &sink,
         .mergeStderrToStdout = true,
-    });
+    }).wait();
 }
 
 void DerivationGoal::buildDone()
diff --git a/src/libutil/processes.cc b/src/libutil/processes.cc
index bd5ccdb4b..3c353d82e 100644
--- a/src/libutil/processes.cc
+++ b/src/libutil/processes.cc
@@ -249,7 +249,7 @@ std::pair<int, std::string> runProgram(RunOptions && options)
     int status = 0;
 
     try {
-        runProgram2(options);
+        runProgram2(options).wait();
     } catch (ExecError & e) {
         status = e.status;
     }
@@ -257,7 +257,25 @@ std::pair<int, std::string> runProgram(RunOptions && options)
     return {status, std::move(sink.s)};
 }
 
-void runProgram2(const RunOptions & options)
+RunningProgram::~RunningProgram()
+{
+    if (pid) {
+        // we will not kill a subprocess because we *can't* kill a subprocess
+        // reliably without placing it in its own process group, and cleaning
+        // up a subprocess only when `separatePG` is set is a loaded footgun.
+        assert(false && "destroying un-wait()ed running process");
+        std::terminate();
+    }
+}
+
+void RunningProgram::wait()
+{
+    int status = pid.wait();
+    if (status)
+        throw ExecError(status, "program '%1%' %2%", program, statusToString(status));
+}
+
+RunningProgram runProgram2(const RunOptions & options)
 {
     checkInterrupt();
 
@@ -317,11 +335,7 @@ void runProgram2(const RunOptions & options)
     if (options.standardOut)
         *options.standardOut << drainFDSource(out.readSide.get());
 
-    /* Wait for the child to finish. */
-    int status = pid.wait();
-
-    if (status)
-        throw ExecError(status, "program '%1%' %2%", options.program, statusToString(status));
+    return RunningProgram{options.program, std::move(pid)};
 }
 
 std::string statusToString(int status)
diff --git a/src/libutil/processes.hh b/src/libutil/processes.hh
index 48a195481..90b2d20af 100644
--- a/src/libutil/processes.hh
+++ b/src/libutil/processes.hh
@@ -87,9 +87,26 @@ struct RunOptions
     bool isInteractive = false;
 };
 
+struct [[nodiscard("you must call RunningProgram::wait()")]] RunningProgram
+{
+    friend RunningProgram runProgram2(const RunOptions & options);
+
+private:
+    Path program;
+    Pid pid;
+
+    RunningProgram(Path program, Pid pid) : program(std::move(program)), pid(std::move(pid)) {}
+
+public:
+    RunningProgram() = default;
+    ~RunningProgram();
+
+    void wait();
+};
+
 std::pair<int, std::string> runProgram(RunOptions && options);
 
-void runProgram2(const RunOptions & options);
+RunningProgram runProgram2(const RunOptions & options);
 
 class ExecError : public Error
 {