Support: Add polling option to sys::Wait

[arsenm]

Summary

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

WIP because of windows support. No idea if the Windows part even
builds, much less works.

[arsenm]

Rebase on top of std::optional change

[aganea]

llvm/lib/Support/Unix/Program.inc

447

Out of curiosity (sorry I am not well versed into Linux APIs) -- who does stop the child process & when?

[arsenm]

llvm/lib/Support/Unix/Program.inc

447

Above a signal handler was installed. The process receives SIGALRM after the timeout from the OS. The original path here then killed off the child; this instead sends the signal to continue

[aganea]

llvm/lib/Support/Unix/Program.inc

447

The alarm signal/handler is received in the context of the parent process. The child process' state is unaffected by the alarm, right? I suppose the alarm has the effect to make wait4() return (from what I understand from the doc), but again, this doesn't affect the state of the child, does it? The child should be in a 'running' state at this point. So my question is, why do we need to send SIGCONT at all to the child?

[arsenm]

Don't send SIGCONT, rename test

llvm/lib/Support/Unix/Program.inc

447

I don't know what I was doing before, but I added this to fix it waiting forever. Maybe I just had another bug in an earlier revision, my tests work with this removed

[aganea]

llvm/lib/Support/Unix/Program.inc

428

I wouldn't return right away, ProcStat is not yet updated at this point. I think enclosing all this into if (!Polling) like before would be fine. Seems good otherwise.

[arsenm]

Address comment. Build bot also implies it at least builds on windows with a passed unit test

[aganea]

LGTM.

[arsenm]

eb93b8774e2cfc5172f0727db8a50427adcb69d7

llvm/include/llvm/Support/Program.h

#endif
 
   /// This function waits for the process specified by \p PI to finish.
   /// \returns A \see ProcessInfo struct with Pid set to:
   /// \li The process id of the child process if the child process has changed
   /// state.
   /// \li 0 if the child process has not changed state.
   /// \note Users of this function should always check the ReturnCode member of
   /// the \see ProcessInfo returned from this function.
-  ProcessInfo Wait(
-      const ProcessInfo &PI, ///< The child process that should be waited on.
+  ProcessInfo
+  Wait(const ProcessInfo &PI, ///< The child process that should be waited on.
        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.
   void printArg(llvm::raw_ostream &OS, StringRef Arg, bool Quote);
 
 #if defined(_WIN32)
   /// Given a list of command line arguments, quote and escape them as necessary
   /// to build a single flat command line appropriate for calling CreateProcess
   /// on
   /// Windows.
   ErrorOr<std::wstring> flattenWindowsCommandLine(ArrayRef<StringRef> Args);
 #endif
   }
 }
 
 #endif

llvm/lib/Support/Unix/Program.inc

pid_t(llvm::sys::wait4)(pid_t pid, int *status, int options,
   // absence of indiscriminate `wait` calls from the current process (which
   // would cause the call here to fail with ECHILD).
   return ::wait4(pid, status, options & ~WNOHANG, usage);
 }
 #endif
 
 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?");
 
   int WaitPidOptions = 0;
   pid_t ChildPid = PI.Pid;
   bool WaitUntilTerminates = false;
   if (!SecondsToWait) {
     WaitUntilTerminates = true;
   } else if (*SecondsToWait != 0) {
     // Install a timeout handler.  The handler itself does nothing, but the
     // simple fact of having a handler at all causes the wait below to return
     // with EINTR, unlike if we used SIG_IGN.
     memset(&Act, 0, sizeof(Act));
     Act.sa_handler = TimeOutHandler;
     sigemptyset(&Act.sa_mask);
     sigaction(SIGALRM, &Act, &Old);
     // FIXME The alarm signal may be delivered to another thread.
-
     alarm(*SecondsToWait);
   } else
     WaitPidOptions = WNOHANG;
 
   // Parent process: Wait for the child process to terminate.
   int status;
   ProcessInfo WaitResult;
   rusage Info;
   if (ProcStat)
     ProcStat->reset();
 
   do {
     WaitResult.Pid = sys::wait4(ChildPid, &status, WaitPidOptions, &Info);
   } while (WaitUntilTerminates && WaitResult.Pid == -1 && errno == EINTR);
 
   if (WaitResult.Pid != PI.Pid) {
     if (WaitResult.Pid == 0) {
       // Non-blocking wait.
       return WaitResult;
     } else {
       if (SecondsToWait && errno == EINTR) {
+        if (!Polling) {
           // Kill the child.
           kill(PI.Pid, SIGKILL);
 
           // Turn off the alarm and restore the signal handler
           alarm(0);
           sigaction(SIGALRM, &Old, nullptr);
 
           // Wait for child to die
           // FIXME This could grab some other child process out from another
           // waiting thread and then leave a zombie anyway.
           if (wait(&status) != ChildPid)
             MakeErrMsg(ErrMsg, "Child timed out but wouldn't die");
           else
             MakeErrMsg(ErrMsg, "Child timed out", 0);
 
           WaitResult.ReturnCode = -2; // Timeout detected
+        } else {
+          // Child's not done yet, resume it.
+          kill(PI.Pid, SIGCONT);
+        }
+
         return WaitResult;
       } else if (errno != EINTR) {
         MakeErrMsg(ErrMsg, "Error waiting for child process");
         WaitResult.ReturnCode = -1;
         return WaitResult;
       }
     }
   }

llvm/lib/Support/Windows/Program.inc

ErrorOr<std::wstring> sys::flattenWindowsCommandLine(ArrayRef<StringRef> Args) {
   SmallVector<wchar_t, MAX_PATH> CommandUtf16;
   if (std::error_code ec = windows::UTF8ToUTF16(Command, CommandUtf16))
     return ec;
 
   return std::wstring(CommandUtf16.begin(), CommandUtf16.end());
 }
 
 ProcessInfo sys::Wait(const ProcessInfo &PI, std::optional<unsigned> SecondsToWait,
-                      std::string *ErrMsg, std::optional<ProcessStatistics> *ProcStat) {
+                      std::string *ErrMsg, 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?");
   DWORD milliSecondsToWait = SecondsToWait ? *SecondsToWait * 1000 : INFINITE;
 
   ProcessInfo WaitResult = PI;
   if (ProcStat)
     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");
 
         // -2 indicates a crash or timeout as opposed to failure to execute.
         WaitResult.ReturnCode = -2;
         CloseHandle(PI.Process);
         return WaitResult;

llvm/unittests/Support/ProgramTest.cpp

addEnvVar("LLVM_PROGRAM_TEST_TIMEOUT=1");
   std::string Error;
   bool ExecutionFailed;
   int RetCode =
       ExecuteAndWait(Executable, argv, getEnviron(), {}, /*SecondsToWait*/ 1,
                      /*MemoryLimit*/ 0, &Error, &ExecutionFailed);
   ASSERT_EQ(-2, RetCode);
 }
 
+TEST_F(ProgramEnvTest, TestExecuteNoWaitTimeoutNoKill) {
+  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.TestExecuteNoWaitTimeoutNoKill"};
+
+  // 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};
 
   {
     std::string Error;
     bool ExecutionFailed;
     int RetCode = ExecuteAndWait(Executable, argv, std::nullopt, {}, 0, 0,