Extend Command for more control over a single exectuion. #Centipede

This refines the original Execute() into ExecuteAsync() + Wait(), with RequestStop() to interrupt the execution.

This allows Centipede to properly handle Command timeout (moved to CentipedeCallbacks::RunBatchForBinary) no matter if it is using a fork server or not.

PiperOrigin-RevId: 781159715

Author: xinhaoyuan

centipede/centipede_callbacks.cc

@@ -18,6 +18,8 @@
 #include <cstddef>
 #include <cstdlib>
 #include <filesystem>  // NOLINT
+#include <memory>
+#include <optional>
 #include <string>
 #include <string_view>
 #include <system_error>  // NOLINT
@@ -146,7 +148,7 @@ std::string CentipedeCallbacks::ConstructRunnerFlags(
 Command &CentipedeCallbacks::GetOrCreateCommandForBinary(
     std::string_view binary) {
   for (auto &cmd : commands_) {
-    if (cmd.path() == binary) return cmd;
+    if (cmd->path() == binary) return *cmd;
   }
   // We don't want to collect coverage for extra binaries. It won't be used.
   bool disable_coverage =
@@ -165,25 +167,51 @@ Command &CentipedeCallbacks::GetOrCreateCommandForBinary(
         absl::StrCat("LLVM_PROFILE_FILE=",
                      WorkDir{env_}.SourceBasedCoverageRawProfilePath()));
 
-  // Allow for the time it takes to fork a subprocess etc.
-  const auto amortized_timeout =
-      env_.timeout_per_batch == 0
-          ? absl::InfiniteDuration()
-          : absl::Seconds(env_.timeout_per_batch) + absl::Seconds(5);
   Command::Options cmd_options;
   cmd_options.env_add = std::move(env);
   cmd_options.env_remove = EnvironmentVariablesToUnset();
   cmd_options.stdout_file = execute_log_path_;
   cmd_options.stderr_file = execute_log_path_;
-  cmd_options.timeout = amortized_timeout;
   cmd_options.temp_file_path = temp_input_file_path_;
-  Command &cmd =
-      commands_.emplace_back(Command{binary, std::move(cmd_options)});
+  Command &cmd = *commands_.emplace_back(
+      std::make_unique<Command>(binary, std::move(cmd_options)));
   if (env_.fork_server) cmd.StartForkServer(temp_dir_, Hash(binary));
 
   return cmd;
 }
 
+int CentipedeCallbacks::RunBatchForBinary(std::string_view binary) {
+  auto &cmd = GetOrCreateCommandForBinary(binary);
+  const absl::Duration amortized_timeout =
+      env_.timeout_per_batch == 0
+          ? absl::InfiniteDuration()
+          : absl::Seconds(env_.timeout_per_batch) + absl::Seconds(5);
+  const auto deadline = absl::Now() + amortized_timeout;
+  int exit_code = EXIT_SUCCESS;
+  const bool should_clean_up = [&] {
+    if (!cmd.ExecuteAsync()) return true;
+    const std::optional<int> ret = cmd.Wait(deadline);
+    if (!ret.has_value()) return true;
+    exit_code = *ret;
+    return false;
+  }();
+  if (should_clean_up) {
+    exit_code = [&] {
+      if (!cmd.is_executing()) return EXIT_FAILURE;
+      LOG(ERROR) << "Cleaning up the batch execution.";
+      cmd.RequestStop();
+      const auto ret = cmd.Wait(absl::Now() + absl::Seconds(60));
+      if (ret.has_value()) return *ret;
+      LOG(ERROR) << "Batch execution cleanup failed to end in 60s.";
+      return EXIT_FAILURE;
+    }();
+    commands_.erase(
+        std::find_if(commands_.begin(), commands_.end(),
+                     [=](const auto &cmd) { return cmd->path() == binary; }));
+  }
+  return exit_code;
+}
+
 int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
     std::string_view binary, const std::vector<ByteArray> &inputs,
     BatchResult &batch_result) {
@@ -212,12 +240,11 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
   }
 
   // Run.
-  Command &cmd = GetOrCreateCommandForBinary(binary);
-  int retval = cmd.Execute();
+  const int exit_code = RunBatchForBinary(binary);
   inputs_blobseq_.ReleaseSharedMemory();  // Inputs are already consumed.
 
   // Get results.
-  batch_result.exit_code() = retval;
+  batch_result.exit_code() = exit_code;
   const bool read_success = batch_result.Read(outputs_blobseq_);
   LOG_IF(ERROR, !read_success) << "Failed to read batch result!";
   outputs_blobseq_.ReleaseSharedMemory();  // Outputs are already consumed.
@@ -229,7 +256,7 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
   //   * Will be logged by the caller.
   // * some outputs were not written because the outputs_blobseq_ overflown.
   //   * Logged by the following code.
-  if (retval == 0 && read_success &&
+  if (exit_code == 0 && read_success &&
       batch_result.num_outputs_read() != num_inputs_written) {
     LOG(INFO) << "Read " << batch_result.num_outputs_read() << "/"
               << num_inputs_written
@@ -239,7 +266,7 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
 
   if (env_.print_runner_log) PrintExecutionLog();
 
-  if (retval != EXIT_SUCCESS) {
+  if (exit_code != EXIT_SUCCESS) {
     ReadFromLocalFile(execute_log_path_, batch_result.log());
     ReadFromLocalFile(failure_description_path_,
                       batch_result.failure_description());
@@ -257,7 +284,7 @@ int CentipedeCallbacks::ExecuteCentipedeSancovBinaryWithShmem(
     std::filesystem::remove(failure_signature_path_);
   }
   VLOG(1) << __FUNCTION__ << " took " << (absl::Now() - start_time);
-  return retval;
+  return exit_code;
 }
 
 // See also: `DumpSeedsToDir()`.
@@ -366,19 +393,18 @@ MutationResult CentipedeCallbacks::MutateViaExternalBinary(
       << VV(num_inputs_written) << VV(inputs.size());
 
   // Execute.
-  Command &cmd = GetOrCreateCommandForBinary(binary);
-  int retval = cmd.Execute();
+  const int exit_code = RunBatchForBinary(binary);
   inputs_blobseq_.ReleaseSharedMemory();  // Inputs are already consumed.
 
-  if (retval != EXIT_SUCCESS) {
-    LOG(WARNING) << "Custom mutator failed with exit code: " << retval;
+  if (exit_code != EXIT_SUCCESS) {
+    LOG(WARNING) << "Custom mutator failed with exit code: " << exit_code;
   }
-  if (env_.print_runner_log || retval != EXIT_SUCCESS) {
+  if (env_.print_runner_log || exit_code != EXIT_SUCCESS) {
     PrintExecutionLog();
   }
 
   MutationResult result;
-  result.exit_code() = retval;
+  result.exit_code() = exit_code;
   result.Read(num_mutants, outputs_blobseq_);
   outputs_blobseq_.ReleaseSharedMemory();  // Outputs are already consumed.
 

centipede/centipede_callbacks.h

@@ -161,6 +161,8 @@ class CentipedeCallbacks {
   // Returns a Command object with matching `binary` from commands_,
   // creates one if needed.
   Command &GetOrCreateCommandForBinary(std::string_view binary);
+  // Runs a batch with the command `binary` and returns the exit code.
+  int RunBatchForBinary(std::string_view binary);
 
   // Prints the execution log from the last executed binary.
   void PrintExecutionLog() const;
@@ -182,7 +184,8 @@ class CentipedeCallbacks {
   SharedMemoryBlobSequence inputs_blobseq_;
   SharedMemoryBlobSequence outputs_blobseq_;
 
-  std::vector<Command> commands_;
+  // Need unique_ptr indirection because Command is not movable/copyable.
+  std::vector<std::unique_ptr<Command>> commands_;
 };
 
 // Abstract class for creating/destroying CentipedeCallbacks objects.

centipede/command.cc

@@ -16,9 +16,11 @@
 
 #include <errno.h>
 #include <fcntl.h>
+#include <spawn.h>
 #include <sys/poll.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <sys/wait.h>
 #include <unistd.h>
 #ifdef __APPLE__
 #include <inttypes.h>
@@ -30,6 +32,7 @@
 #include <cstdlib>
 #include <filesystem>  // NOLINT
 #include <fstream>
+#include <optional>
 #include <string>
 #include <string_view>
 #include <system_error>  // NOLINT
@@ -55,6 +58,12 @@
 #include "./centipede/util.h"
 #include "./common/logging.h"
 
+#if !defined(_MSC_VER)
+// Needed to pass the current environment to posix_spawn, which needs an
+// explicit envp without an option to inherit implicitly.
+extern char **environ;
+#endif
+
 namespace fuzztest::internal {
 namespace {
 
@@ -141,8 +150,16 @@ struct Command::ForkServerProps {
 // the deleter is instantiated, the special member functions must be defined
 // out-of-line here, now that ForkServerProps is complete (that's by-the-book
 // PIMPL).
-Command::Command(Command &&other) noexcept = default;
-Command::~Command() = default;
+Command::~Command() {
+  if (is_executing()) {
+    LOG(WARNING)
+        << "Destructing Command object for " << path() << " with "
+        << (fork_server_ ? absl::StrCat("fork server PID ", fork_server_->pid_)
+                         : absl::StrCat("PID ", pid_))
+        << " still running. Requesting it to stop without waiting for it...";
+    RequestStop();
+  }
+}
 
 Command::Command(std::string_view path, Options options)
     : path_(path), options_(std::move(options)) {}
@@ -308,88 +325,101 @@ absl::Status Command::VerifyForkServerIsHealthy() {
   return absl::OkStatus();
 }
 
-int Command::Execute() {
+bool Command::ExecuteAsync() {
+  CHECK(!is_executing());
   VLOG(1) << "Executing command '" << command_line_ << "'...";
 
-  int exit_code = EXIT_SUCCESS;
-
   if (fork_server_ != nullptr) {
-    VLOG(1) << "Sending execution request to fork server: "
-            << VV(options_.timeout);
+    VLOG(1) << "Sending execution request to fork server";
 
     if (const auto status = VerifyForkServerIsHealthy(); !status.ok()) {
       LogProblemInfo(absl::StrCat("Fork server should be running, but isn't: ",
                                   status.message()));
-      return EXIT_FAILURE;
+      return false;
     }
 
     // Wake up the fork server.
     char x = ' ';
     CHECK_EQ(1, write(fork_server_->pipe_[0], &x, 1));
+  } else {
+    CHECK_EQ(pid_, -1);
+    std::vector<std::string> argv_strs = {"/bin/sh", "-c", command_line_};
+    std::vector<char *> argv;
+    argv.reserve(argv_strs.size() + 1);
+    for (auto &argv_str : argv_strs) {
+      argv.push_back(argv_str.data());
+    }
+    argv.push_back(nullptr);
+    CHECK_EQ(posix_spawn(&pid_, argv[0], /*file_actions=*/nullptr,
+                         /*attrp=*/nullptr, argv.data(), environ),
+             0);
+  }
+
+  is_executing_ = true;
+  return true;
+}
 
+std::optional<int> Command::Wait(absl::Time deadline) {
+  CHECK(is_executing());
+  int exit_code = EXIT_SUCCESS;
+
+  if (fork_server_ != nullptr) {
     // The fork server forks, the child is running. Block until some readable
     // data appears in the pipe (that is, after the fork server writes the
     // execution result to it).
     struct pollfd poll_fd = {};
     int poll_ret = -1;
-    auto poll_deadline = absl::Now() + options_.timeout;
-    bool sigterm_sent = false;
-    bool try_again = false;
     do {
-      try_again = false;
       // NOTE: `poll_fd` has to be reset every time.
       poll_fd = {
           /*fd=*/fork_server_->pipe_[1],  // The file descriptor to wait for.
           /*events=*/POLLIN,              // Wait until `fd` gets readable data.
       };
       const int poll_timeout_ms = static_cast<int>(absl::ToInt64Milliseconds(
-          std::max(poll_deadline - absl::Now(), absl::Milliseconds(1))));
+          std::max(deadline - absl::Now(), absl::Milliseconds(1))));
       poll_ret = poll(&poll_fd, 1, poll_timeout_ms);
       // The `poll()` syscall can get interrupted: it sets errno==EINTR in that
       // case. We should tolerate that.
-      if (poll_ret < 0 && errno == EINTR) {
-        try_again = true;
-        continue;
-      }
-      if (poll_ret == 0 && !sigterm_sent) {
-        LogProblemInfo(
-            absl::StrCat("Timeout while waiting for fork server: timeout is ",
-                         absl::FormatDuration(options_.timeout)));
-        CHECK_NE(fork_server_->pid_, -1);
-        LOG(INFO) << "Sending SIGTERM to the fork server PID "
-                  << fork_server_->pid_ << " and waiting for 60s";
-        kill(fork_server_->pid_, SIGTERM);
-        sigterm_sent = true;
-        poll_deadline += absl::Seconds(60);
-        try_again = true;
-        continue;
-      }
-    } while (try_again);
-
+    } while (poll_ret < 0 && errno == EINTR);
     if (poll_ret != 1 || (poll_fd.revents & POLLIN) == 0) {
       // The fork server errored out or timed out, or some other error occurred,
       // e.g. the syscall was interrupted.
       if (poll_ret == 0) {
-        CHECK(sigterm_sent);
-        LogProblemInfo(
-            "Fork server did not respond within 60s after SIGTERM was sent");
-        // TODO: xinhaoyuan - the right thing to do is to either properly
-        // recover or request early exit.
+        LogProblemInfo(absl::StrCat(
+            "Timeout while waiting for fork server: deadline is ", deadline));
       } else {
         LogProblemInfo(absl::StrCat(
             "Error while waiting for fork server: poll() returned ", poll_ret));
       }
-      return EXIT_FAILURE;
+      return std::nullopt;
     }
 
     // The fork server wrote the execution result to the pipe: read it.
     CHECK_EQ(sizeof(exit_code),
              read(fork_server_->pipe_[1], &exit_code, sizeof(exit_code)));
   } else {
-    VLOG(1) << "Fork server disabled - executing command directly";
-    // No fork server, use system().
-    exit_code = system(command_line_.c_str());
+    CHECK_NE(pid_, -1);
+    while (true) {
+      const pid_t r = waitpid(pid_, &exit_code, WNOHANG);
+      CHECK_NE(r, -1);
+      if (r == pid_ && (WIFEXITED(exit_code) || WIFSIGNALED(exit_code))) break;
+      CHECK_EQ(r, 0);
+      const auto timeout = deadline - absl::Now();
+      if (timeout > absl::ZeroDuration()) {
+        const auto duration = std::clamp<useconds_t>(
+            absl::ToInt64Microseconds(timeout), 0, 100000);
+        usleep(duration);  // NOLINT: early return on SIGCHLD is desired.
+        continue;
+      } else {
+        LogProblemInfo(absl::StrCat(
+            "Timeout while waiting for the command process: deadline is ",
+            deadline));
+        return std::nullopt;
+      }
+    }
+    pid_ = -1;
   }
+  is_executing_ = false;
 
   // When the command is actually a wrapper shell launching the binary(-es)
   // (e.g. a Docker container), the shell will preserve a normal exit code
@@ -444,6 +474,17 @@ int Command::Execute() {
   return exit_code;
 }
 
+void Command::RequestStop() {
+  CHECK(is_executing());
+  if (fork_server_) {
+    CHECK_NE(fork_server_->pid_, -1);
+    kill(fork_server_->pid_, SIGTERM);
+    return;
+  }
+  CHECK_NE(pid_, -1);
+  kill(pid_, SIGTERM);
+}
+
 std::string Command::ReadRedirectedStdout() const {
   std::string ret;
   if (!options_.stdout_file.empty()) {