Support: Add polling option to sys::Wait

Currently the process is terminated after the timeout. Add an option
to let the process resume after the timeout instead.

https://reviews.llvm.org/D138952

llvm/include/llvm/Support/Program.h

@@ -210,17 +210,22 @@ namespace sys {
        std::optional<unsigned> SecondsToWait, ///< If std::nullopt, waits until
        ///< child has terminated.
        ///< If a value, this specifies the amount of time to wait for the child
-       ///< process to exit. If the time expires, the child is killed and this
-       ///< function returns. If zero, this function will perform a non-blocking
+       ///< process. If the time expires, and \p Polling is false, the child is
+       ///< killed and this < function returns. If the time expires and \p
+       ///< Polling is true, the child is resumed.
+       ///<
+       ///< If zero, this function will perform a non-blocking
        ///< wait on the child process.
        std::string *ErrMsg = nullptr, ///< If non-zero, provides a pointer to a
        ///< string instance in which error messages will be returned. If the
        ///< string is non-empty upon return an error occurred while invoking the
        ///< program.
        std::optional<ProcessStatistics> *ProcStat =
-           nullptr ///< If non-zero, provides
+           nullptr, ///< If non-zero, provides
        /// a pointer to a structure in which process execution statistics will
        /// be stored.
+
+       bool Polling = false ///< If true, do not kill the process on timeout.
   );
 
   /// Print a command argument, and optionally quote it.

llvm/lib/Support/Unix/Program.inc

@@ -387,7 +387,8 @@ pid_t(llvm::sys::wait4)(pid_t pid, int *status, int options,
 ProcessInfo llvm::sys::Wait(const ProcessInfo &PI,
                             std::optional<unsigned> SecondsToWait,
                             std::string *ErrMsg,
-                            std::optional<ProcessStatistics> *ProcStat) {
+                            std::optional<ProcessStatistics> *ProcStat,
+                            bool Polling) {
   struct sigaction Act, Old;
   assert(PI.Pid && "invalid pid to wait on, process not started?");
 
@@ -408,12 +409,11 @@ ProcessInfo llvm::sys::Wait(const ProcessInfo &PI,
     sigemptyset(&Act.sa_mask);
     sigaction(SIGALRM, &Act, &Old);
     // FIXME The alarm signal may be delivered to another thread.
-
     alarm(*SecondsToWait);
   }
 
   // Parent process: Wait for the child process to terminate.
-  int status;
+  int status = 0;
   ProcessInfo WaitResult;
   rusage Info;
   if (ProcStat)
@@ -428,7 +428,7 @@ ProcessInfo llvm::sys::Wait(const ProcessInfo &PI,
       // Non-blocking wait.
       return WaitResult;
     } else {
-      if (SecondsToWait && errno == EINTR) {
+      if (SecondsToWait && errno == EINTR && !Polling) {
         // Kill the child.
         kill(PI.Pid, SIGKILL);
 

llvm/lib/Support/Windows/Program.inc

@@ -411,7 +411,8 @@ ErrorOr<std::wstring> sys::flattenWindowsCommandLine(ArrayRef<StringRef> Args) {
 ProcessInfo sys::Wait(const ProcessInfo &PI,
                       std::optional<unsigned> SecondsToWait,
                       std::string *ErrMsg,
-                      std::optional<ProcessStatistics> *ProcStat) {
+                      std::optional<ProcessStatistics> *ProcStat,
+                      bool Polling) {
   assert(PI.Pid && "invalid pid to wait on, process not started?");
   assert((PI.Process && PI.Process != INVALID_HANDLE_VALUE) &&
          "invalid process handle to wait on, process not started?");
@@ -422,7 +423,7 @@ ProcessInfo sys::Wait(const ProcessInfo &PI,
     ProcStat->reset();
   DWORD WaitStatus = WaitForSingleObject(PI.Process, milliSecondsToWait);
   if (WaitStatus == WAIT_TIMEOUT) {
-    if (*SecondsToWait > 0) {
+    if (!Polling && *SecondsToWait > 0) {
       if (!TerminateProcess(PI.Process, 1)) {
         if (ErrMsg)
           MakeErrMsg(ErrMsg, "Failed to terminate timed-out program");

llvm/unittests/Support/ProgramTest.cpp

@@ -284,6 +284,45 @@ TEST_F(ProgramEnvTest, TestExecuteAndWaitTimeout) {
   ASSERT_EQ(-2, RetCode);
 }
 
+TEST_F(ProgramEnvTest, TestExecuteNoWaitTimeoutPolling) {
+  using namespace llvm::sys;
+
+  if (getenv("LLVM_PROGRAM_TEST_TIMEOUT")) {
+    sleep_for(/*seconds*/ 5);
+    exit(0);
+  }
+
+  std::string Executable =
+      sys::fs::getMainExecutable(TestMainArgv0, &ProgramTestStringArg1);
+  StringRef argv[] = {
+      Executable,
+      "--gtest_filter=ProgramEnvTest.TestExecuteNoWaitTimeoutPolling"};
+
+  // Add LLVM_PROGRAM_TEST_TIMEOUT to the environment of the child.
+  addEnvVar("LLVM_PROGRAM_TEST_TIMEOUT=1");
+
+  std::string Error;
+  bool ExecutionFailed;
+  ProcessInfo PI0 = ExecuteNoWait(Executable, argv, getEnviron(),
+                                  /*Redirects=*/{}, /*MemoryLimit=*/0, &Error,
+                                  &ExecutionFailed);
+  ASSERT_FALSE(ExecutionFailed) << Error;
+  ASSERT_NE(PI0.Pid, ProcessInfo::InvalidPid) << "Invalid process id";
+
+  // Check that we don't kill the process with a non-0 SecondsToWait if Polling.
+  unsigned LoopCount = 0;
+  ProcessInfo WaitResult;
+  do {
+    ++LoopCount;
+    WaitResult = llvm::sys::Wait(PI0, /*SecondsToWait=*/1, &Error,
+                                 /*ProcStats=*/nullptr,
+                                 /*Polling=*/true);
+    ASSERT_TRUE(Error.empty()) << Error;
+  } while (WaitResult.Pid != PI0.Pid);
+
+  ASSERT_GT(LoopCount, 1u) << "LoopCount should be >1";
+}
+
 TEST(ProgramTest, TestExecuteNegative) {
   std::string Executable = "i_dont_exist";
   StringRef argv[] = {Executable};