rework the Subprocess::communicate() API
[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  * Base exception thrown by the Subprocess methods.
149  */
150 class SubprocessError : public std::exception {};
151
152 /**
153  * Exception thrown by *Checked methods of Subprocess.
154  */
155 class CalledProcessError : public SubprocessError {
156  public:
157   explicit CalledProcessError(ProcessReturnCode rc);
158   ~CalledProcessError() throw() { }
159   const char* what() const throw() FOLLY_OVERRIDE { return what_.c_str(); }
160   ProcessReturnCode returnCode() const { return returnCode_; }
161  private:
162   ProcessReturnCode returnCode_;
163   std::string what_;
164 };
165
166 /**
167  * Exception thrown if the subprocess cannot be started.
168  */
169 class SubprocessSpawnError : public SubprocessError {
170  public:
171   SubprocessSpawnError(const char* executable, int errCode, int errnoValue);
172   ~SubprocessSpawnError() throw() {}
173   const char* what() const throw() FOLLY_OVERRIDE { return what_.c_str(); }
174   int errnoValue() const { return errnoValue_; }
175
176  private:
177   int errnoValue_;
178   std::string what_;
179 };
180
181 /**
182  * Subprocess.
183  */
184 class Subprocess : private boost::noncopyable {
185  public:
186   static const int CLOSE = -1;
187   static const int PIPE = -2;
188   static const int PIPE_IN = -3;
189   static const int PIPE_OUT = -4;
190
191   /**
192    * Class representing various options: file descriptor behavior, and
193    * whether to use $PATH for searching for the executable,
194    *
195    * By default, we don't use $PATH, file descriptors are closed if
196    * the close-on-exec flag is set (fcntl FD_CLOEXEC) and inherited
197    * otherwise.
198    */
199   class Options : private boost::orable<Options> {
200     friend class Subprocess;
201    public:
202     Options()
203       : closeOtherFds_(false),
204         usePath_(false),
205         parentDeathSignal_(0) {
206     }
207
208     /**
209      * Change action for file descriptor fd.
210      *
211      * "action" may be another file descriptor number (dup2()ed before the
212      * child execs), or one of CLOSE, PIPE_IN, and PIPE_OUT.
213      *
214      * CLOSE: close the file descriptor in the child
215      * PIPE_IN: open a pipe *from* the child
216      * PIPE_OUT: open a pipe *to* the child
217      *
218      * PIPE is a shortcut; same as PIPE_IN for stdin (fd 0), same as
219      * PIPE_OUT for stdout (fd 1) or stderr (fd 2), and an error for
220      * other file descriptors.
221      */
222     Options& fd(int fd, int action);
223
224     /**
225      * Shortcut to change the action for standard input.
226      */
227     Options& stdin(int action) { return fd(STDIN_FILENO, action); }
228
229     /**
230      * Shortcut to change the action for standard output.
231      */
232     Options& stdout(int action) { return fd(STDOUT_FILENO, action); }
233
234     /**
235      * Shortcut to change the action for standard error.
236      * Note that stderr(1) will redirect the standard error to the same
237      * file descriptor as standard output; the equivalent of bash's "2>&1"
238      */
239     Options& stderr(int action) { return fd(STDERR_FILENO, action); }
240
241     Options& pipeStdin() { return fd(STDIN_FILENO, PIPE_IN); }
242     Options& pipeStdout() { return fd(STDOUT_FILENO, PIPE_OUT); }
243     Options& pipeStderr() { return fd(STDERR_FILENO, PIPE_OUT); }
244
245     /**
246      * Close all other fds (other than standard input, output, error,
247      * and file descriptors explicitly specified with fd()).
248      *
249      * This is potentially slow; it's generally a better idea to
250      * set the close-on-exec flag on all file descriptors that shouldn't
251      * be inherited by the child.
252      *
253      * Even with this option set, standard input, output, and error are
254      * not closed; use stdin(CLOSE), stdout(CLOSE), stderr(CLOSE) if you
255      * desire this.
256      */
257     Options& closeOtherFds() { closeOtherFds_ = true; return *this; }
258
259     /**
260      * Use the search path ($PATH) when searching for the executable.
261      */
262     Options& usePath() { usePath_ = true; return *this; }
263
264     /**
265      * Child will receive a signal when the parent exits.
266      */
267     Options& parentDeathSignal(int sig) {
268       parentDeathSignal_ = sig;
269       return *this;
270     }
271
272     /**
273      * Helpful way to combine Options.
274      */
275     Options& operator|=(const Options& other);
276
277    private:
278     typedef boost::container::flat_map<int, int> FdMap;
279     FdMap fdActions_;
280     bool closeOtherFds_;
281     bool usePath_;
282     int parentDeathSignal_;
283   };
284
285   static Options pipeStdin() { return Options().stdin(PIPE); }
286   static Options pipeStdout() { return Options().stdout(PIPE); }
287   static Options pipeStderr() { return Options().stderr(PIPE); }
288
289   /**
290    * Create a subprocess from the given arguments.  argv[0] must be listed.
291    * If not-null, executable must be the actual executable
292    * being used (otherwise it's the same as argv[0]).
293    *
294    * If env is not-null, it must contain name=value strings to be used
295    * as the child's environment; otherwise, we inherit the environment
296    * from the parent.  env must be null if options.usePath is set.
297    */
298   explicit Subprocess(
299       const std::vector<std::string>& argv,
300       const Options& options = Options(),
301       const char* executable = nullptr,
302       const std::vector<std::string>* env = nullptr);
303   ~Subprocess();
304
305   /**
306    * Create a subprocess run as a shell command (as shell -c 'command')
307    *
308    * The shell to use is taken from the environment variable $SHELL,
309    * or /bin/sh if $SHELL is unset.
310    */
311   explicit Subprocess(
312       const std::string& cmd,
313       const Options& options = Options(),
314       const std::vector<std::string>* env = nullptr);
315
316   /**
317    * Communicate with the child until all pipes to/from the child are closed.
318    *
319    * The input buffer is written to the process' stdin pipe, and data is read
320    * from the stdout and stderr pipes.  Non-blocking I/O is performed on all
321    * pipes simultaneously to avoid deadlocks.
322    *
323    * The stdin pipe will be closed after the full input buffer has been written.
324    * An error will be thrown if a non-empty input buffer is supplied but stdin
325    * was not configured as a pipe.
326    *
327    * Returns a pair of buffers containing the data read from stdout and stderr.
328    * If stdout or stderr is not a pipe, an empty IOBuf queue will be returned
329    * for the respective buffer.
330    *
331    * Note that communicate() returns when all pipes to/from the child are
332    * closed; the child might stay alive after that, so you must still wait().
333    *
334    * communicateIOBuf uses IOBufQueue for buffering (which has the advantage
335    * that it won't try to allocate all data at once).  communicate
336    * uses strings for simplicity.
337    */
338   std::pair<IOBufQueue, IOBufQueue> communicateIOBuf(
339       IOBufQueue input = IOBufQueue());
340
341   std::pair<std::string, std::string> communicate(
342       StringPiece input = StringPiece());
343
344   /**
345    * Communicate with the child until all pipes to/from the child are closed.
346    *
347    * readCallback(pfd, cfd) will be called whenever there's data available
348    * on any pipe *from* the child (PIPE_OUT).  pfd is the file descriptor
349    * in the parent (that you use to read from); cfd is the file descriptor
350    * in the child (used for identifying the stream; 1 = child's standard
351    * output, 2 = child's standard error, etc)
352    *
353    * writeCallback(pfd, cfd) will be called whenever a pipe *to* the child is
354    * writable (PIPE_IN).  pfd is the file descriptor in the parent (that you
355    * use to write to); cfd is the file descriptor in the child (used for
356    * identifying the stream; 0 = child's standard input, etc)
357    *
358    * The read and write callbacks must read from / write to pfd and return
359    * false during normal operation or true at end-of-file;
360    * communicate() will then close the pipe.  Note that pfd is
361    * nonblocking, so be prepared for read() / write() to return -1 and
362    * set errno to EAGAIN (in which case you should return false).
363    *
364    * NOTE that you MUST consume all data passed to readCallback (or return
365    * true, which will close the pipe, possibly sending SIGPIPE to the child or
366    * making its writes fail with EPIPE), and you MUST write to a writable pipe
367    * (or return true, which will close the pipe).  To do otherwise is an
368    * error.  You must do this even for pipes you are not interested in.
369    *
370    * Note that communicate() returns when all pipes to/from the child are
371    * closed; the child might stay alive after that, so you must still wait().
372    *
373    * Most users won't need to use this; the simpler version of communicate
374    * (which buffers data in memory) will probably work fine.
375    */
376   typedef std::function<bool(int, int)> FdCallback;
377   void communicate(FdCallback readCallback, FdCallback writeCallback);
378
379   /**
380    * Return the child's pid, or -1 if the child wasn't successfully spawned
381    * or has already been wait()ed upon.
382    */
383   pid_t pid() const;
384
385   /**
386    * Return the child's status (as per wait()) if the process has already
387    * been waited on, -1 if the process is still running, or -2 if the process
388    * hasn't been successfully started.  NOTE that this does not poll, but
389    * returns the status stored in the Subprocess object.
390    */
391   ProcessReturnCode returnCode() const { return returnCode_; }
392
393   /**
394    * Poll the child's status and return it, return -1 if the process
395    * is still running.  NOTE that it is illegal to call poll again after
396    * poll indicated that the process has terminated, or to call poll on a
397    * process that hasn't been successfully started (the constructor threw an
398    * exception).
399    */
400   ProcessReturnCode poll();
401
402   /**
403    * Poll the child's status.  If the process is still running, return false.
404    * Otherwise, return true if the process exited with status 0 (success),
405    * or throw CalledProcessError if the process exited with a non-zero status.
406    */
407   bool pollChecked();
408
409   /**
410    * Wait for the process to terminate and return its status.
411    * Similarly to poll, it is illegal to call wait after the process
412    * has already been reaped or if the process has not successfully started.
413    */
414   ProcessReturnCode wait();
415
416   /**
417    * Wait for the process to terminate, throw if unsuccessful.
418    */
419   void waitChecked();
420
421   /**
422    * Set all pipes from / to child non-blocking.  communicate() does
423    * this for you.
424    */
425   void setAllNonBlocking();
426
427   /**
428    * Get parent file descriptor corresponding to the given file descriptor
429    * in the child.  Throws if childFd isn't a pipe (PIPE_IN / PIPE_OUT).
430    * Do not close() the return file descriptor; use closeParentFd, below.
431    */
432   int parentFd(int childFd) const {
433     return pipes_[findByChildFd(childFd)].parentFd;
434   }
435   int stdin() const { return parentFd(0); }
436   int stdout() const { return parentFd(1); }
437   int stderr() const { return parentFd(2); }
438
439   /**
440    * Close the parent file descriptor given a file descriptor in the child.
441    */
442   void closeParentFd(int childFd);
443
444   /**
445    * Send a signal to the child.  Shortcuts for the commonly used Unix
446    * signals are below.
447    */
448   void sendSignal(int signal);
449   void terminate() { sendSignal(SIGTERM); }
450   void kill() { sendSignal(SIGKILL); }
451
452  private:
453   static const int RV_RUNNING = ProcessReturnCode::RV_RUNNING;
454   static const int RV_NOT_STARTED = ProcessReturnCode::RV_NOT_STARTED;
455
456   // spawn() sets up a pipe to read errors from the child,
457   // then calls spawnInternal() to do the bulk of the work.  Once
458   // spawnInternal() returns it reads the error pipe to see if the child
459   // encountered any errors.
460   void spawn(
461       std::unique_ptr<const char*[]> argv,
462       const char* executable,
463       const Options& options,
464       const std::vector<std::string>* env);
465   void spawnInternal(
466       std::unique_ptr<const char*[]> argv,
467       const char* executable,
468       Options& options,
469       const std::vector<std::string>* env,
470       int errFd);
471
472   // Actions to run in child.
473   // Note that this runs after vfork(), so tread lightly.
474   // Returns 0 on success, or an errno value on failure.
475   int prepareChild(const Options& options, const sigset_t* sigmask) const;
476   int runChild(const char* executable, char** argv, char** env,
477                const Options& options) const;
478
479   /**
480    * Read from the error pipe, and throw SubprocessSpawnError if the child
481    * failed before calling exec().
482    */
483   void readChildErrorPipe(int pfd, const char* executable);
484
485   /**
486    * Close all file descriptors.
487    */
488   void closeAll();
489
490   // return index in pipes_
491   int findByChildFd(int childFd) const;
492
493   pid_t pid_;
494   ProcessReturnCode returnCode_;
495
496   // The number of pipes between parent and child is assumed to be small,
497   // so we're happy with a vector here, even if it means linear erase.
498   // sorted by childFd
499   struct PipeInfo : private boost::totally_ordered<PipeInfo> {
500     int parentFd;
501     int childFd;
502     int direction;  // one of PIPE_IN / PIPE_OUT
503     bool operator<(const PipeInfo& other) const {
504       return childFd < other.childFd;
505     }
506     bool operator==(const PipeInfo& other) const {
507       return childFd == other.childFd;
508     }
509   };
510   std::vector<PipeInfo> pipes_;
511 };
512
513 inline Subprocess::Options& Subprocess::Options::operator|=(
514     const Subprocess::Options& other) {
515   if (this == &other) return *this;
516   // Replace
517   for (auto& p : other.fdActions_) {
518     fdActions_[p.first] = p.second;
519   }
520   closeOtherFds_ |= other.closeOtherFds_;
521   usePath_ |= other.usePath_;
522   return *this;
523 }
524
525 }  // namespace folly
526
527 #endif /* FOLLY_SUBPROCESS_H_ */
528