add Subprocess::waitOrTerminateOrKill

Summary: `waitOrTerminateOrKill` fills the gap between `wait()` and `terminateOrKill`. It allows giving the subprocess a chance to shut down cleanly (after closing its stdin pipe, for example), but avoids waiting forever or leaking a child process.

Reviewed By: yfeldblum, mzhaom

Differential Revision: D17769090

fbshipit-source-id: 8940fd63f7eb9c09ef293c89a5e97f69805735ff

folly/Subprocess.cpp

@@ -666,23 +666,17 @@ void Subprocess::waitChecked() {
   checkStatus(returnCode_);
 }
 
-void Subprocess::sendSignal(int signal) {
-  returnCode_.enforce(ProcessReturnCode::RUNNING);
-  int r = ::kill(pid_, signal);
-  checkUnixError(r, "kill");
-}
-
-ProcessReturnCode Subprocess::terminateOrKill(TimeoutDuration sigtermTimeout) {
+ProcessReturnCode Subprocess::waitTimeout(TimeoutDuration timeout) {
   returnCode_.enforce(ProcessReturnCode::RUNNING);
   DCHECK_GT(pid_, 0) << "The subprocess has been waited already";
-  // 1. Send SIGTERM to kill the process
-  terminate();
-  // 2. check whether subprocess has terminated using non-blocking waitpid
-  auto pollUntil = std::chrono::steady_clock::now() + sigtermTimeout;
-  do {
+
+  auto pollUntil = std::chrono::steady_clock::now() + timeout;
+  for (;;) {
+    // Always call waitpid once after the full timeout has elapsed.
+    auto now = std::chrono::steady_clock::now();
+
     int status;
     pid_t found;
-    // warp waitpid in the while loop to deal with EINTR.
     do {
       found = ::waitpid(pid_, &status, WNOHANG);
     } while (found == -1 && errno == EINTR);
@@ -696,10 +690,45 @@ ProcessReturnCode Subprocess::terminateOrKill(TimeoutDuration sigtermTimeout) {
       pid_ = -1;
       return returnCode_;
     }
-
+    if (now > pollUntil) {
+      // Timed out: still running().
+      return returnCode_;
+    }
     // The subprocess is still running, sleep for 100ms
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
-  } while (std::chrono::steady_clock::now() <= pollUntil);
+    std::this_thread::sleep_for(std::chrono::milliseconds{100});
+  }
+}
+
+void Subprocess::sendSignal(int signal) {
+  returnCode_.enforce(ProcessReturnCode::RUNNING);
+  int r = ::kill(pid_, signal);
+  checkUnixError(r, "kill");
+}
+
+ProcessReturnCode Subprocess::waitOrTerminateOrKill(
+    TimeoutDuration waitTimeout,
+    TimeoutDuration sigtermTimeout) {
+  returnCode_.enforce(ProcessReturnCode::RUNNING);
+  DCHECK_GT(pid_, 0) << "The subprocess has been waited already";
+
+  this->waitTimeout(waitTimeout);
+
+  if (returnCode_.running()) {
+    return terminateOrKill(sigtermTimeout);
+  }
+  return returnCode_;
+}
+
+ProcessReturnCode Subprocess::terminateOrKill(TimeoutDuration sigtermTimeout) {
+  returnCode_.enforce(ProcessReturnCode::RUNNING);
+  DCHECK_GT(pid_, 0) << "The subprocess has been waited already";
+  // 1. Send SIGTERM to kill the process
+  terminate();
+  // 2. check whether subprocess has terminated using non-blocking waitpid
+  waitTimeout(sigtermTimeout);
+  if (!returnCode_.running()) {
+    return returnCode_;
+  }
   // 3. If we are at this point, we have waited enough time after
   // sending SIGTERM, we have to use nuclear option SIGKILL to kill
   // the subprocess.

folly/Subprocess.h

@@ -593,7 +593,7 @@ class Subprocess {
   /**
    * Wait for the process to terminate and return its status.  Like poll(),
    * the only exception this can throw is std::logic_error if you call this
-   * on a Subprocess whose status is RUNNING.  Aborts on egregious
+   * on a Subprocess whose status is not RUNNING.  Aborts on egregious
    * violations of contract, like an out-of-band waitpid(p.pid(), 0, 0).
    */
   ProcessReturnCode wait();
@@ -603,6 +603,16 @@ class Subprocess {
    */
   void waitChecked();
 
+  using TimeoutDuration = std::chrono::milliseconds;
+
+  /**
+   * Call `waitpid` non-blockingly up to `timeout`. Throws std::logic_error if
+   * called on a Subprocess whose status is not RUNNING.
+   *
+   * The return code will be running() if waiting timed out.
+   */
+  ProcessReturnCode waitTimeout(TimeoutDuration timeout);
+
   /**
    * Send a signal to the child.  Shortcuts for the commonly used Unix
    * signals are below.
@@ -615,7 +625,14 @@ class Subprocess {
     sendSignal(SIGKILL);
   }
 
-  using TimeoutDuration = std::chrono::milliseconds;
+  /**
+   * Call `waitpid` non-blockingly up to `waitTimeout`. If the process hasn't
+   * terminated after that, fall back on `terminateOrKill` with
+   * `sigtermTimeoutSeconds`.
+   */
+  ProcessReturnCode waitOrTerminateOrKill(
+      TimeoutDuration waitTimeout,
+      TimeoutDuration sigtermTimeout);
 
   /**
    * Send the SIGTERM to terminate the process, poll `waitpid` non-blockingly

folly/test/SubprocessTest.cpp

@@ -174,6 +174,20 @@ TEST(SimpleSubprocessTest, ChangeChildDirectoryWithError) {
   }
 }
 
+TEST(SimpleSubprocessTest, waitOrTerminateOrKill_waits_if_process_exits) {
+  Subprocess proc(std::vector<std::string>{"/bin/sleep", "0.1"});
+  auto retCode = proc.waitOrTerminateOrKill(1s, 1s);
+  EXPECT_TRUE(retCode.exited());
+  EXPECT_EQ(0, retCode.exitStatus());
+}
+
+TEST(SimpleSubprocessTest, waitOrTerminateOrKill_terminates_if_timeout) {
+  Subprocess proc(std::vector<std::string>{"/bin/sleep", "60"});
+  auto retCode = proc.waitOrTerminateOrKill(1s, 1s);
+  EXPECT_TRUE(retCode.killed());
+  EXPECT_EQ(SIGTERM, retCode.killSignal());
+}
+
 // This method verifies terminateOrKill shouldn't affect the exit
 // status if the process has exitted already.
 TEST(SimpleSubprocessTest, TerminateAfterProcessExit) {