git.m6w6.name Git - m6w6/libmemcached/blob - test/lib/ForkAndExec.cpp

   1 #include "ForkAndExec.hpp"
   2 #include "p9y/poll.hpp"
   3
   4 #include <cstdio>
   5 #include <fcntl.h>
   6 #if HAVE_UNISTD_H
   7 #  include <unistd.h>
   8 #endif
   9 #include <spawn.h>
  10
  11 #if !HAVE_PIPE2
  12 static inline int setfl(int fd, int newflags) {
  13   auto oldflags = fcntl(fd, F_GETFL);
  14   return fcntl(fd, F_SETFL, oldflags | newflags);
  15 }
  16 static inline int pipe2(int pipefd[2], int flags) {
  17   int rc;
  18
  19   rc = pipe(pipefd);
  20   if (0 > rc) {
  21     return rc;
  22   }
  23   rc = setfl(pipefd[0], flags);
  24   if (0 > rc) {
  25     return rc;
  26   }
  27   return setfl(pipefd[1], flags);
  28 }
  29 #endif
  30
  31 ForkAndExec::ForkAndExec(const char *binary_, char **argv_)
  32 : ready{-1, -1}
  33 , pipes{-1, -1}
  34 , binary{binary_}
  35 , argv{argv_}
  36 {
  37 }
  38
  39 ForkAndExec::~ForkAndExec() {
  40   for (auto &pipe : {ready, pipes}) {
  41     for (auto m : {READ, WRITE}) {
  42       closePipe(pipe[m]);
  43     }
  44   }
  45 }
  46
  47 int ForkAndExec::createPipe() {
  48   if (pipes[mode::READ] == -1) {
  49     if(pipe2(pipes, O_NONBLOCK)) {
  50       perror("ForkAndExec pipe()");
  51       return -1;
  52     }
  53   }
  54
  55   return pipes[mode::READ];
  56 }
  57
  58 pid_t ForkAndExec::operator()()  {
  59   if (!prepareExecReadyPipe()) {
  60     return 0;
  61   }
  62
  63   switch (pid_t pid = fork()) {
  64     case 0:
  65       closePipe(pipes[mode::READ]);
  66       prepareOutputPipe();
  67       execvp(binary, argv);
  68       perror("ForkAndExec exec()");
  69       _exit(EXIT_FAILURE);
  70
  71     case -1:
  72       perror("ForkAndExec fork()");
  73       return 0;
  74
  75     default:
  76       pipes[mode::READ] = -1;
  77       closePipe(pipes[mode::WRITE]);
  78       pollExecReadyPipe();
  79       return pid;
  80   }
  81 }
  82
  83 bool ForkAndExec::prepareExecReadyPipe() {
  84   if (ready[mode::READ] == -1) {
  85     if (pipe2(ready, O_CLOEXEC | O_NONBLOCK)) {
  86       perror("ForkAndExec pipe()");
  87       return false;
  88     }
  89     closePipe(ready[mode::WRITE]);
  90   }
  91
  92   return true;
  93 }
  94
  95 void ForkAndExec::prepareOutputPipe() {
  96   if (pipes[mode::WRITE] != -1) {
  97     if (0 > dup2(pipes[mode::WRITE], STDERR_FILENO) ||
  98         0 > dup2(pipes[mode::WRITE], STDOUT_FILENO)) {
  99       perror("ForkAndExec dup()");
 100     }
 101   }
 102 }
 103
 104 void ForkAndExec::closePipe(int &fd) {
 105   if (fd != -1) {
 106     close(fd);
 107     fd = -1;
 108   }
 109 }
 110
 111 void ForkAndExec::pollExecReadyPipe() {
 112 #if __APPLE__
 113   char c, n = 50;
 114   do {
 115     if (0 == read(ready[mode::READ], &c, 1)) {
 116       return;
 117     }
 118     this_thread::sleep_for(100ms);
 119   } while (errno == EAGAIN && n--);
 120 #else
 121   pollfd fd{ready[mode::READ], 0, 0};
 122   if (1 > poll(&fd, 1, 5000)) {
 123     cerr << "exec() timed out" << endl;
 124   }
 125 #endif
 126 }