common/fork_function: helper to run a function, forked, with a timeout

Signed-off-by: Sage Weil <sage@redhat.com>

diff --git a/src/common/fork_function.h b/src/common/fork_function.h
new file mode 100644 (file)
index 0000000..56b8a3b
--- /dev/null
+++ b/src/common/fork_function.h
@@ -0,0 +1,146 @@
+// -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:t -*-
+// vim: ts=8 sw=2 smarttab
+
+// Run a function in a forked child, with a timeout.
+
+#pragma once
+
+#include <functional>
+#include <signal.h>
+#include <sys/wait.h>
+#include <sys/types.h>
+#include <common/errno.h>
+#include <ostream>
+#include "common/errno.h"
+
+static void _fork_function_dummy_sighandler(int sig) {}
+
+// Run a function post-fork, with a timeout.  Function can return
+// int8_t only due to unix exit code limitations.  Returns -ETIMEDOUT
+// if timeout is reached.
+
+static inline int fork_function(
+  int timeout,
+  std::ostream& errstr,
+  std::function<int8_t(void)> f)
+{
+  // first fork the forker.
+  pid_t forker_pid = fork();
+  if (forker_pid) {
+    // just wait
+    int status;
+    while (waitpid(forker_pid, &status, 0) == -1) {
+      assert(errno == EINTR);
+    }
+    if (WIFSIGNALED(status)) {
+      errstr << ": got signal: " << WTERMSIG(status) << "\n";
+      return 128 + WTERMSIG(status);
+    }
+    if (WIFEXITED(status)) {
+      int8_t r = WEXITSTATUS(status);
+      errstr << ": exit status: " << (int)r << "\n";
+      return r;
+    }
+    errstr << ": waitpid: unknown status returned\n";
+    return -1;
+  }
+
+  // we are forker (first child)
+  sigset_t mask, oldmask;
+  int pid;
+
+  // Restore default action for SIGTERM in case the parent process decided
+  // to ignore it.
+  if (signal(SIGTERM, SIG_DFL) == SIG_ERR) {
+    std::cerr << ": signal failed: " << cpp_strerror(errno) << "\n";
+    goto fail_exit;
+  }
+  // Because SIGCHLD is ignored by default, setup dummy handler for it,
+  // so we can mask it.
+  if (signal(SIGCHLD, _fork_function_dummy_sighandler) == SIG_ERR) {
+    std::cerr << ": signal failed: " << cpp_strerror(errno) << "\n";
+    goto fail_exit;
+  }
+  // Setup timeout handler.
+  if (signal(SIGALRM, timeout_sighandler) == SIG_ERR) {
+    std::cerr << ": signal failed: " << cpp_strerror(errno) << "\n";
+    goto fail_exit;
+  }
+  // Block interesting signals.
+  sigemptyset(&mask);
+  sigaddset(&mask, SIGINT);
+  sigaddset(&mask, SIGTERM);
+  sigaddset(&mask, SIGCHLD);
+  sigaddset(&mask, SIGALRM);
+  if (sigprocmask(SIG_SETMASK, &mask, &oldmask) == -1) {
+    std::cerr << ": sigprocmask failed: "
+             << cpp_strerror(errno) << "\n";
+    goto fail_exit;
+  }
+
+  pid = fork();
+
+  if (pid == -1) {
+    std::cerr << ": fork failed: " << cpp_strerror(errno) << "\n";
+    goto fail_exit;
+  }
+
+  if (pid == 0) { // we are second child
+    // Restore old sigmask.
+    if (sigprocmask(SIG_SETMASK, &oldmask, NULL) == -1) {
+      std::cerr << ": sigprocmask failed: "
+               << cpp_strerror(errno) << "\n";
+      goto fail_exit;
+    }
+    (void)setpgid(0, 0); // Become process group leader.
+    int8_t r = f();
+    _exit((uint8_t)r);
+  }
+
+  // Parent
+  (void)alarm(timeout);
+
+  for (;;) {
+    int signo;
+    if (sigwait(&mask, &signo) == -1) {
+      std::cerr << ": sigwait failed: " << cpp_strerror(errno) << "\n";
+      goto fail_exit;
+    }
+    switch (signo) {
+    case SIGCHLD:
+      int status;
+      if (waitpid(pid, &status, WNOHANG) == -1) {
+       std::cerr << ": waitpid failed: " << cpp_strerror(errno) << "\n";
+       goto fail_exit;
+      }
+      if (WIFEXITED(status))
+       _exit(WEXITSTATUS(status));
+      if (WIFSIGNALED(status))
+       _exit(128 + WTERMSIG(status));
+      std::cerr << ": unknown status returned\n";
+      goto fail_exit;
+    case SIGINT:
+    case SIGTERM:
+      // Pass SIGINT and SIGTERM, which are usually used to terminate
+      // a process, to the child.
+      if (::kill(pid, signo) == -1) {
+       std::cerr << ": kill failed: " << cpp_strerror(errno) << "\n";
+       goto fail_exit;
+      }
+      continue;
+    case SIGALRM:
+      std::cerr << ": timed out (" << timeout << " sec)\n";
+      if (::killpg(pid, SIGKILL) == -1) {
+       std::cerr << ": kill failed: " << cpp_strerror(errno) << "\n";
+       goto fail_exit;
+      }
+      _exit(-ETIMEDOUT);
+    default:
+      std::cerr << ": sigwait: invalid signal: " << signo << "\n";
+      goto fail_exit;
+    }
+  }
+  return 0;
+fail_exit:
+  _exit(EXIT_FAILURE);
+}

diff --git a/src/test/test_subprocess.cc b/src/test/test_subprocess.cc
index 29c7b94acd7f4f521e9f29d883da698f08124117..c316eb58a129c9327006de62454ccc58a518baff 100644 (file)
--- a/src/test/test_subprocess.cc
+++ b/src/test/test_subprocess.cc
@@ -21,6 +21,7 @@
 #include "common/SubProcess.h"
 #include "common/safe_io.h"
 #include "gtest/gtest.h"
+#include "common/fork_function.h"
 
 bool read_from_fd(int fd, std::string &out) {
   out.clear();
@@ -266,3 +267,27 @@ TEST(SubProcessTimed, SubshellTimedout)
   std::cerr << "err: " << sh.err() << std::endl;
   ASSERT_FALSE(sh.err().c_str()[0] == '\0');
 }
+
+TEST(fork_function, normal)
+{
+  ASSERT_EQ(0, fork_function(10, std::cerr, [&]() { return 0; }));
+  ASSERT_EQ(1, fork_function(10, std::cerr, [&]() { return 1; }));
+  ASSERT_EQ(13, fork_function(10, std::cerr, [&]() { return 13; }));
+  ASSERT_EQ(-1, fork_function(10, std::cerr, [&]() { return -1; }));
+  ASSERT_EQ(-13, fork_function(10, std::cerr, [&]() { return -13; }));
+  ASSERT_EQ(-ETIMEDOUT,
+           fork_function(10, std::cerr, [&]() { return -ETIMEDOUT; }));
+}
+
+TEST(fork_function, timeout)
+{
+  ASSERT_EQ(-ETIMEDOUT, fork_function(2, std::cerr, [&]() {
+       sleep(60);
+       return 0; }));
+  ASSERT_EQ(-ETIMEDOUT, fork_function(2, std::cerr, [&]() {
+       sleep(60);
+       return 1; }));
+  ASSERT_EQ(-ETIMEDOUT, fork_function(2, std::cerr, [&]() {
+       sleep(60);
+       return -111; }));
+}