graduate IOBuf out of folly/experimental
[folly.git] / folly / Subprocess.h
1 /*
2  * Copyright 2013 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *   http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 /**
18  * Subprocess library, modeled after Python's subprocess module
19  * (http://docs.python.org/2/library/subprocess.html)
20  *
21  * This library defines one class (Subprocess) which represents a child
22  * process.  Subprocess has two constructors: one that takes a vector<string>
23  * and executes the given executable without using the shell, and one
24  * that takes a string and executes the given command using the shell.
25  * Subprocess allows you to redirect the child's standard input, standard
26  * output, and standard error to/from child descriptors in the parent,
27  * or to create communication pipes between the child and the parent.
28  *
29  * The simplest example is a thread-safe version of the system() library
30  * function:
31  *    Subprocess(cmd).wait();
32  * which executes the command using the default shell and waits for it
33  * to complete, returning the exit status.
34  *
35  * A thread-safe version of popen() (type="r", to read from the child):
36  *    Subprocess proc(cmd, Subprocess::pipeStdout());
37  *    // read from proc.stdout()
38  *    proc.wait();
39  *
40  * A thread-safe version of popen() (type="w", to write from the child):
41  *    Subprocess proc(cmd, Subprocess::pipeStdin());
42  *    // write to proc.stdin()
43  *    proc.wait();
44  *
45  * If you want to redirect both stdin and stdout to pipes, you can, but
46  * note that you're subject to a variety of deadlocks.  You'll want to use
47  * nonblocking I/O; look at the implementation of communicate() for an example.
48  *
49  * communicate() is a way to communicate to a child via its standard input,
50  * standard output, and standard error.  It buffers everything in memory,
51  * so it's not great for large amounts of data (or long-running processes),
52  * but it insulates you from the deadlocks mentioned above.
53  */
54 #ifndef FOLLY_SUBPROCESS_H_
55 #define FOLLY_SUBPROCESS_H_
56
57 #include <sys/types.h>
58 #include <signal.h>
59 #include <wait.h>
60
61 #include <exception>
62 #include <vector>
63 #include <string>
64
65 #include <boost/container/flat_map.hpp>
66 #include <boost/operators.hpp>
67 #include <boost/noncopyable.hpp>
68
69 #include "folly/io/IOBufQueue.h"
70 #include "folly/MapUtil.h"
71 #include "folly/Portability.h"
72 #include "folly/Range.h"
73
74 namespace folly {
75
76 /**
77  * Class to wrap a process return code.
78  */
79 class Subprocess;
80 class ProcessReturnCode {
81   friend class Subprocess;
82  public:
83   enum State {
84     NOT_STARTED,
85     RUNNING,
86     EXITED,
87     KILLED
88   };
89
90   /**
91    * Process state.  One of:
92    * NOT_STARTED: process hasn't been started successfully
93    * RUNNING: process is currently running
94    * EXITED: process exited (successfully or not)
95    * KILLED: process was killed by a signal.
96    */
97   State state() const;
98
99   /**
100    * Helper wrappers around state().
101    */
102   bool notStarted() const { return state() == NOT_STARTED; }
103   bool running() const { return state() == RUNNING; }
104   bool exited() const { return state() == EXITED; }
105   bool killed() const { return state() == KILLED; }
106
107   /**
108    * Exit status.  Only valid if state() == EXITED; throws otherwise.
109    */
110   int exitStatus() const;
111
112   /**
113    * Signal that caused the process's termination.  Only valid if
114    * state() == KILLED; throws otherwise.
115    */
116   int killSignal() const;
117
118   /**
119    * Was a core file generated?  Only valid if state() == KILLED; throws
120    * otherwise.
121    */
122   bool coreDumped() const;
123
124   /**
125    * String representation; one of
126    * "not started"
127    * "running"
128    * "exited with status <status>"
129    * "killed by signal <signal>"
130    * "killed by signal <signal> (core dumped)"
131    */
132   std::string str() const;
133
134   /**
135    * Helper function to enforce a precondition based on this.
136    * Throws std::logic_error if in an unexpected state.
137    */
138   void enforce(State state) const;
139  private:
140   explicit ProcessReturnCode(int rv) : rawStatus_(rv) { }
141   static constexpr int RV_NOT_STARTED = -2;
142   static constexpr int RV_RUNNING = -1;
143
144   int rawStatus_;
145 };
146
147 /**
148  * Exception thrown by *Checked methods of Subprocess.
149  */
150 class CalledProcessError : public std::exception {
151  public:
152   explicit CalledProcessError(ProcessReturnCode rc);
153   ~CalledProcessError() throw() { }
154   const char* what() const throw() FOLLY_OVERRIDE { return what_.c_str(); }
155   ProcessReturnCode returnCode() const { return returnCode_; }
156  private:
157   ProcessReturnCode returnCode_;
158   std::string what_;
159 };
160
161 /**
162  * Subprocess.
163  */
164 class Subprocess : private boost::noncopyable {
165  public:
166   static const int CLOSE = -1;
167   static const int PIPE = -2;
168   static const int PIPE_IN = -3;
169   static const int PIPE_OUT = -4;
170
171   /**
172    * Class representing various options: file descriptor behavior, and
173    * whether to use $PATH for searching for the executable,
174    *
175    * By default, we don't use $PATH, file descriptors are closed if
176    * the close-on-exec flag is set (fcntl FD_CLOEXEC) and inherited
177    * otherwise.
178    */
179   class Options : private boost::orable<Options> {
180     friend class Subprocess;
181    public:
182     Options() : closeOtherFds_(false), usePath_(false) { }
183
184     /**
185      * Change action for file descriptor fd.
186      *
187      * "action" may be another file descriptor number (dup2()ed before the
188      * child execs), or one of CLOSE, PIPE_IN, and PIPE_OUT.
189      *
190      * CLOSE: close the file descriptor in the child
191      * PIPE_IN: open a pipe *from* the child
192      * PIPE_OUT: open a pipe *to* the child
193      *
194      * PIPE is a shortcut; same as PIPE_IN for stdin (fd 0), same as
195      * PIPE_OUT for stdout (fd 1) or stderr (fd 2), and an error for
196      * other file descriptors.
197      */
198     Options& fd(int fd, int action);
199
200     /**
201      * Shortcut to change the action for standard input.
202      */
203     Options& stdin(int action) { return fd(0, action); }
204
205     /**
206      * Shortcut to change the action for standard output.
207      */
208     Options& stdout(int action) { return fd(1, action); }
209
210     /**
211      * Shortcut to change the action for standard error.
212      * Note that stderr(1) will redirect the standard error to the same
213      * file descriptor as standard output; the equivalent of bash's "2>&1"
214      */
215     Options& stderr(int action) { return fd(2, action); }
216
217     /**
218      * Close all other fds (other than standard input, output, error,
219      * and file descriptors explicitly specified with fd()).
220      *
221      * This is potentially slow; it's generally a better idea to
222      * set the close-on-exec flag on all file descriptors that shouldn't
223      * be inherited by the child.
224      *
225      * Even with this option set, standard input, output, and error are
226      * not closed; use stdin(CLOSE), stdout(CLOSE), stderr(CLOSE) if you
227      * desire this.
228      */
229     Options& closeOtherFds() { closeOtherFds_ = true; return *this; }
230
231     /**
232      * Use the search path ($PATH) when searching for the executable.
233      */
234     Options& usePath() { usePath_ = true; return *this; }
235
236     /**
237      * Helpful way to combine Options.
238      */
239     Options& operator|=(const Options& other);
240
241    private:
242     typedef boost::container::flat_map<int, int> FdMap;
243     FdMap fdActions_;
244     bool closeOtherFds_;
245     bool usePath_;
246   };
247
248   static Options pipeStdin() { return Options().stdin(PIPE); }
249   static Options pipeStdout() { return Options().stdout(PIPE); }
250   static Options pipeStderr() { return Options().stderr(PIPE); }
251
252   /**
253    * Create a subprocess from the given arguments.  argv[0] must be listed.
254    * If not-null, executable must be the actual executable
255    * being used (otherwise it's the same as argv[0]).
256    *
257    * If env is not-null, it must contain name=value strings to be used
258    * as the child's environment; otherwise, we inherit the environment
259    * from the parent.  env must be null if options.usePath is set.
260    */
261   explicit Subprocess(
262       const std::vector<std::string>& argv,
263       const Options& options = Options(),
264       const char* executable = nullptr,
265       const std::vector<std::string>* env = nullptr);
266   ~Subprocess();
267
268   /**
269    * Create a subprocess run as a shell command (as shell -c 'command')
270    *
271    * The shell to use is taken from the environment variable $SHELL,
272    * or /bin/sh if $SHELL is unset.
273    */
274   explicit Subprocess(
275       const std::string& cmd,
276       const Options& options = Options(),
277       const std::vector<std::string>* env = nullptr);
278
279   /**
280    * Append all data, close the stdin (to-child) fd, and read all data,
281    * except that this is done in a safe manner to prevent deadlocking.
282    *
283    * If writeStdin() is given in flags, the process must have been opened with
284    * stdinFd=PIPE.
285    *
286    * If readStdout() is given in flags, the first returned value will be the
287    * value read from the child's stdout; the child must have been opened with
288    * stdoutFd=PIPE.
289    *
290    * If readStderr() is given in flags, the second returned value will be the
291    * value read from the child's stderr; the child must have been opened with
292    * stderrFd=PIPE.
293    *
294    * Note that communicate() returns when all pipes to/from the child are
295    * closed; the child might stay alive after that, so you must still wait().
296    *
297    * communicateIOBuf uses IOBufQueue for buffering (which has the advantage
298    * that it won't try to allocate all data at once).  communicate
299    * uses strings for simplicity.
300    */
301   class CommunicateFlags : private boost::orable<CommunicateFlags> {
302     friend class Subprocess;
303    public:
304     CommunicateFlags()
305       : writeStdin_(false), readStdout_(false), readStderr_(false) { }
306     CommunicateFlags& writeStdin() { writeStdin_ = true; return *this; }
307     CommunicateFlags& readStdout() { readStdout_ = true; return *this; }
308     CommunicateFlags& readStderr() { readStderr_ = true; return *this; }
309
310     CommunicateFlags& operator|=(const CommunicateFlags& other);
311    private:
312     bool writeStdin_;
313     bool readStdout_;
314     bool readStderr_;
315   };
316
317   static CommunicateFlags writeStdin() {
318     return CommunicateFlags().writeStdin();
319   }
320   static CommunicateFlags readStdout() {
321     return CommunicateFlags().readStdout();
322   }
323   static CommunicateFlags readStderr() {
324     return CommunicateFlags().readStderr();
325   }
326
327   std::pair<IOBufQueue, IOBufQueue> communicateIOBuf(
328       const CommunicateFlags& flags = readStdout(),
329       IOBufQueue data = IOBufQueue());
330
331   std::pair<std::string, std::string> communicate(
332       const CommunicateFlags& flags = readStdout(),
333       StringPiece data = StringPiece());
334
335   /**
336    * Communicate with the child until all pipes to/from the child are closed.
337    *
338    * readCallback(pfd, cfd) will be called whenever there's data available
339    * on any pipe *from* the child (PIPE_OUT).  pfd is the file descriptor
340    * in the parent (that you use to read from); cfd is the file descriptor
341    * in the child (used for identifying the stream; 1 = child's standard
342    * output, 2 = child's standard error, etc)
343    *
344    * writeCallback(pfd, cfd) will be called whenever a pipe *to* the child is
345    * writable (PIPE_IN).  pfd is the file descriptor in the parent (that you
346    * use to write to); cfd is the file descriptor in the child (used for
347    * identifying the stream; 0 = child's standard input, etc)
348    *
349    * The read and write callbacks must read from / write to pfd and return
350    * false during normal operation or true at end-of-file;
351    * communicate() will then close the pipe.  Note that pfd is
352    * nonblocking, so be prepared for read() / write() to return -1 and
353    * set errno to EAGAIN (in which case you should return false).
354    *
355    * NOTE that you MUST consume all data passed to readCallback (or return
356    * true, which will close the pipe, possibly sending SIGPIPE to the child or
357    * making its writes fail with EPIPE), and you MUST write to a writable pipe
358    * (or return true, which will close the pipe).  To do otherwise is an
359    * error.  You must do this even for pipes you are not interested in.
360    *
361    * Note that communicate() returns when all pipes to/from the child are
362    * closed; the child might stay alive after that, so you must still wait().
363    *
364    * Most users won't need to use this; the simpler version of communicate
365    * (which buffers data in memory) will probably work fine.
366    */
367   typedef std::function<bool(int, int)> FdCallback;
368   void communicate(FdCallback readCallback, FdCallback writeCallback);
369
370   /**
371    * Return the child's pid, or -1 if the child wasn't successfully spawned
372    * or has already been wait()ed upon.
373    */
374   pid_t pid() const;
375
376   /**
377    * Return the child's status (as per wait()) if the process has already
378    * been waited on, -1 if the process is still running, or -2 if the process
379    * hasn't been successfully started.  NOTE that this does not poll, but
380    * returns the status stored in the Subprocess object.
381    */
382   ProcessReturnCode returnCode() const { return returnCode_; }
383
384   /**
385    * Poll the child's status and return it, return -1 if the process
386    * is still running.  NOTE that it is illegal to call poll again after
387    * poll indicated that the process has terminated, or to call poll on a
388    * process that hasn't been successfully started (the constructor threw an
389    * exception).
390    */
391   ProcessReturnCode poll();
392
393   /**
394    * Poll the child's status.  If the process is still running, return false.
395    * Otherwise, return true if the process exited with status 0 (success),
396    * or throw CalledProcessError if the process exited with a non-zero status.
397    */
398   bool pollChecked();
399
400   /**
401    * Wait for the process to terminate and return its status.
402    * Similarly to poll, it is illegal to call wait after the process
403    * has already been reaped or if the process has not successfully started.
404    */
405   ProcessReturnCode wait();
406
407   /**
408    * Wait for the process to terminate, throw if unsuccessful.
409    */
410   void waitChecked();
411
412   /**
413    * Set all pipes from / to child non-blocking.  communicate() does
414    * this for you.
415    */
416   void setAllNonBlocking();
417
418   /**
419    * Get parent file descriptor corresponding to the given file descriptor
420    * in the child.  Throws if childFd isn't a pipe (PIPE_IN / PIPE_OUT).
421    * Do not close() the return file descriptor; use closeParentFd, below.
422    */
423   int parentFd(int childFd) const {
424     return pipes_[findByChildFd(childFd)].parentFd;
425   }
426   int stdin() const { return parentFd(0); }
427   int stdout() const { return parentFd(1); }
428   int stderr() const { return parentFd(2); }
429
430   /**
431    * Close the parent file descriptor given a file descriptor in the child.
432    */
433   void closeParentFd(int childFd);
434
435   /**
436    * Send a signal to the child.  Shortcuts for the commonly used Unix
437    * signals are below.
438    */
439   void sendSignal(int signal);
440   void terminate() { sendSignal(SIGTERM); }
441   void kill() { sendSignal(SIGKILL); }
442
443  private:
444   static const int RV_RUNNING = ProcessReturnCode::RV_RUNNING;
445   static const int RV_NOT_STARTED = ProcessReturnCode::RV_NOT_STARTED;
446
447   void spawn(
448       std::unique_ptr<const char*[]> argv,
449       const char* executable,
450       const Options& options,
451       const std::vector<std::string>* env);
452
453   // Action to run in child.
454   // Note that this runs after vfork(), so tread lightly.
455   void runChild(const char* executable, char** argv, char** env,
456                 const Options& options) const;
457
458   /**
459    * Close all file descriptors.
460    */
461   void closeAll();
462
463   // return index in pipes_
464   int findByChildFd(int childFd) const;
465
466   pid_t pid_;
467   ProcessReturnCode returnCode_;
468
469   // The number of pipes between parent and child is assumed to be small,
470   // so we're happy with a vector here, even if it means linear erase.
471   // sorted by childFd
472   struct PipeInfo : private boost::totally_ordered<PipeInfo> {
473     int parentFd;
474     int childFd;
475     int direction;  // one of PIPE_IN / PIPE_OUT
476     bool operator<(const PipeInfo& other) const {
477       return childFd < other.childFd;
478     }
479     bool operator==(const PipeInfo& other) const {
480       return childFd == other.childFd;
481     }
482   };
483   std::vector<PipeInfo> pipes_;
484 };
485
486 inline Subprocess::Options& Subprocess::Options::operator|=(
487     const Subprocess::Options& other) {
488   if (this == &other) return *this;
489   // Replace
490   for (auto& p : other.fdActions_) {
491     fdActions_[p.first] = p.second;
492   }
493   closeOtherFds_ |= other.closeOtherFds_;
494   usePath_ |= other.usePath_;
495   return *this;
496 }
497
498 inline Subprocess::CommunicateFlags& Subprocess::CommunicateFlags::operator|=(
499     const Subprocess::CommunicateFlags& other) {
500   if (this == &other) return *this;
501   writeStdin_ |= other.writeStdin_;
502   readStdout_ |= other.readStdout_;
503   readStderr_ |= other.readStderr_;
504   return *this;
505 }
506
507 }  // namespace folly
508
509 #endif /* FOLLY_SUBPROCESS_H_ */
510