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

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