1 #pragma once 2 #include <ext/stdio_filebuf.h> 3 #include <sys/wait.h> 4 #include <torch/csrc/profiler/unwind/unwind_error.h> 5 #include <unistd.h> 6 #include <memory> 7 8 namespace torch::unwind { 9 // helper to open a process with stdin/stdout/stderr streams. 10 struct Communicate { CommunicateCommunicate11 Communicate(const char* command, const char** args) { 12 if (pipe(inpipe_) < 0 || pipe(outpipe_) < 0 || pipe(errpipe_) < 0) { 13 throw UnwindError("pipe() failed"); 14 } 15 pid_t pid = fork(); 16 if (pid < 0) { 17 throw UnwindError("fork() failed"); 18 } else if (pid == 0) { // child process 19 close(inpipe_[1]); 20 close(outpipe_[0]); 21 close(errpipe_[0]); 22 23 dup2(inpipe_[0], STDIN_FILENO); 24 dup2(outpipe_[1], STDOUT_FILENO); 25 dup2(errpipe_[1], STDERR_FILENO); 26 execvp(command, (char* const*)args); 27 throw UnwindError("failed execvp"); 28 } else { // parent process 29 close(inpipe_[0]); 30 close(outpipe_[1]); 31 close(errpipe_[1]); 32 outbuf_.reset( 33 new __gnu_cxx::stdio_filebuf<char>(inpipe_[1], std::ios::out)); 34 inbuf_.reset( 35 new __gnu_cxx::stdio_filebuf<char>(outpipe_[0], std::ios::in)); 36 errbuf_.reset( 37 new __gnu_cxx::stdio_filebuf<char>(errpipe_[0], std::ios::in)); 38 in_.reset(new std::istream(inbuf_.get())); 39 out_.reset(new std::ostream(outbuf_.get())); 40 err_.reset(new std::ostream(errbuf_.get())); 41 } 42 } ~CommunicateCommunicate43 ~Communicate() { 44 close(inpipe_[1]); 45 close(outpipe_[0]); 46 close(errpipe_[0]); 47 } outCommunicate48 std::ostream& out() { 49 return *out_; 50 } errCommunicate51 std::ostream& err() { 52 return *err_; 53 } inCommunicate54 std::istream& in() { 55 return *in_; 56 } 57 58 private: 59 int inpipe_[2]; 60 int outpipe_[2]; 61 int errpipe_[2]; 62 std::unique_ptr<__gnu_cxx::stdio_filebuf<char>> outbuf_, inbuf_, errbuf_; 63 std::unique_ptr<std::istream> in_; 64 std::unique_ptr<std::ostream> out_; 65 std::unique_ptr<std::ostream> err_; 66 }; 67 68 } // namespace torch::unwind 69