coco
538 lines
16 KiB
C++
538 lines
16 KiB
C++
// Copyright (C) 2016 Davis E. King (davis@dlib.net)
|
|
// License: Boost Software License See LICENSE.txt for the full license.
|
|
|
|
#include "subprocess_stream.h"
|
|
|
|
#include <sstream>
|
|
#include <utility>
|
|
#include <iostream>
|
|
#include <cstdio>
|
|
#include <fcntl.h>
|
|
#include <signal.h>
|
|
#include <sys/wait.h>
|
|
#include <sys/select.h>
|
|
#include "call_matlab.h"
|
|
|
|
using namespace std;
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
namespace dlib
|
|
{
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
void make_fd_non_blocking(int fd)
|
|
{
|
|
int flags = fcntl(fd, F_GETFL, 0);
|
|
fcntl(fd, F_SETFL, flags | O_NONBLOCK);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
// Block until fd is ready to read, while also echoing whatever is in fd_printf to
|
|
// cout.
|
|
int read_echoing_select(int fd, int fd_printf)
|
|
{
|
|
// run until fd has data ready
|
|
while(fd_printf >= 0)
|
|
{
|
|
fd_set rfds;
|
|
int retval;
|
|
|
|
while(true)
|
|
{
|
|
FD_ZERO(&rfds);
|
|
FD_SET(fd, &rfds);
|
|
FD_SET(fd_printf, &rfds);
|
|
|
|
// select times out every second just so we can check for matlab ctrl+c.
|
|
struct timeval tv;
|
|
tv.tv_sec = 1;
|
|
tv.tv_usec = 0;
|
|
|
|
try{check_for_matlab_ctrl_c();} catch(...) { return 1; }
|
|
retval = select(std::max(fd,fd_printf)+1, &rfds, NULL, NULL, &tv);
|
|
try{check_for_matlab_ctrl_c();} catch(...) { return 1; }
|
|
if (retval == 0) // keep going if it was just a timeout.
|
|
continue;
|
|
else if (retval == -1 && errno == EINTR)
|
|
continue;
|
|
|
|
break;
|
|
}
|
|
|
|
if (retval == -1)
|
|
{
|
|
return 1;
|
|
}
|
|
else
|
|
{
|
|
if (FD_ISSET(fd,&rfds))
|
|
{
|
|
return 0;
|
|
}
|
|
else
|
|
{
|
|
char buf[1024];
|
|
int num = read(fd_printf,buf, sizeof(buf)-1);
|
|
if (num == -1)
|
|
return 1;
|
|
if (num > 0)
|
|
{
|
|
buf[num] = 0;
|
|
cout << buf << flush;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
int write_echoing_select(int fd, int fd_printf)
|
|
{
|
|
// run until fd has data ready
|
|
while(fd_printf >= 0)
|
|
{
|
|
fd_set rfds, wfds;
|
|
int retval;
|
|
while(true)
|
|
{
|
|
FD_ZERO(&rfds);
|
|
FD_ZERO(&wfds);
|
|
FD_SET(fd, &wfds);
|
|
FD_SET(fd_printf, &rfds);
|
|
|
|
// select times out every second just so we can check for matlab ctrl+c.
|
|
struct timeval tv;
|
|
tv.tv_sec = 1;
|
|
tv.tv_usec = 0;
|
|
|
|
try{check_for_matlab_ctrl_c();} catch(...) { return 1; }
|
|
retval = select(std::max(fd,fd_printf)+1, &rfds, &wfds, NULL, &tv);
|
|
try{check_for_matlab_ctrl_c();} catch(...) { return 1; }
|
|
if (retval == 0) // keep going if it was just a timeout.
|
|
continue;
|
|
else if (retval == -1 && errno == EINTR)
|
|
continue;
|
|
|
|
break;
|
|
}
|
|
|
|
if (retval == -1)
|
|
{
|
|
return 1;
|
|
}
|
|
else
|
|
{
|
|
if (FD_ISSET(fd,&wfds))
|
|
{
|
|
return 0;
|
|
}
|
|
else
|
|
{
|
|
char buf[1024];
|
|
int num = read(fd_printf,buf, sizeof(buf)-1);
|
|
if (num == -1)
|
|
return 1;
|
|
if (num > 0)
|
|
{
|
|
buf[num] = 0;
|
|
cout << buf << flush;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
class filestreambuf : public std::streambuf
|
|
{
|
|
/*!
|
|
INITIAL VALUE
|
|
- fd == the file descriptor we read from.
|
|
- in_buffer == an array of in_buffer_size bytes
|
|
- out_buffer == an array of out_buffer_size bytes
|
|
|
|
CONVENTION
|
|
- in_buffer == the input buffer used by this streambuf
|
|
- out_buffer == the output buffer used by this streambuf
|
|
- max_putback == the maximum number of chars to have in the put back buffer.
|
|
!*/
|
|
|
|
public:
|
|
|
|
filestreambuf (
|
|
int fd_,
|
|
int fd_printf_
|
|
) :
|
|
fd(fd_),
|
|
fd_printf(fd_printf_),
|
|
out_buffer(0),
|
|
in_buffer(0)
|
|
{
|
|
init();
|
|
}
|
|
|
|
virtual ~filestreambuf (
|
|
)
|
|
{
|
|
sync();
|
|
delete [] out_buffer;
|
|
delete [] in_buffer;
|
|
}
|
|
|
|
int sync (
|
|
)
|
|
{
|
|
if (flush_out_buffer() == EOF)
|
|
{
|
|
// an error occurred
|
|
return -1;
|
|
}
|
|
return 0;
|
|
}
|
|
protected:
|
|
|
|
void init (
|
|
)
|
|
{
|
|
try
|
|
{
|
|
out_buffer = new char[out_buffer_size];
|
|
in_buffer = new char[in_buffer_size];
|
|
}
|
|
catch (...)
|
|
{
|
|
if (out_buffer) delete [] out_buffer;
|
|
throw;
|
|
}
|
|
setp(out_buffer, out_buffer + (out_buffer_size-1));
|
|
setg(in_buffer+max_putback,
|
|
in_buffer+max_putback,
|
|
in_buffer+max_putback);
|
|
}
|
|
|
|
int flush_out_buffer (
|
|
)
|
|
{
|
|
int num = static_cast<int>(pptr()-pbase());
|
|
const int num_written = num;
|
|
char* buf = out_buffer;
|
|
while(num != 0)
|
|
{
|
|
if(write_echoing_select(fd, fd_printf))
|
|
return EOF;
|
|
int status = write(fd,buf,num);
|
|
if (status < 0)
|
|
{
|
|
// the write was not successful so return EOF
|
|
return EOF;
|
|
}
|
|
num -= status;
|
|
buf += status;
|
|
}
|
|
pbump(-num_written);
|
|
return num_written;
|
|
}
|
|
|
|
// output functions
|
|
int_type overflow (
|
|
int_type c
|
|
)
|
|
{
|
|
if (c != EOF)
|
|
{
|
|
*pptr() = c;
|
|
pbump(1);
|
|
}
|
|
if (flush_out_buffer() == EOF)
|
|
{
|
|
// an error occurred
|
|
return EOF;
|
|
}
|
|
return c;
|
|
}
|
|
|
|
|
|
std::streamsize xsputn (
|
|
const char* s,
|
|
std::streamsize num
|
|
)
|
|
{
|
|
// Add a sanity check here
|
|
DLIB_ASSERT(num >= 0,
|
|
"\tstd::streamsize filestreambuf::xsputn"
|
|
<< "\n\tThe number of bytes to write can't be negative"
|
|
<< "\n\tnum: " << num
|
|
<< "\n\tthis: " << this
|
|
);
|
|
|
|
std::streamsize space_left = static_cast<std::streamsize>(epptr()-pptr());
|
|
if (num <= space_left)
|
|
{
|
|
std::memcpy(pptr(),s,static_cast<size_t>(num));
|
|
pbump(static_cast<int>(num));
|
|
return num;
|
|
}
|
|
else
|
|
{
|
|
std::memcpy(pptr(),s,static_cast<size_t>(space_left));
|
|
s += space_left;
|
|
pbump(space_left);
|
|
std::streamsize num_left = num - space_left;
|
|
|
|
if (flush_out_buffer() == EOF)
|
|
{
|
|
// the write was not successful so return that 0 bytes were written
|
|
return 0;
|
|
}
|
|
|
|
if (num_left < out_buffer_size)
|
|
{
|
|
std::memcpy(pptr(),s,static_cast<size_t>(num_left));
|
|
pbump(num_left);
|
|
return num;
|
|
}
|
|
else
|
|
{
|
|
while(num_left != 0)
|
|
{
|
|
if(write_echoing_select(fd, fd_printf))
|
|
return EOF;
|
|
int status = write(fd,s,num_left);
|
|
if (status < 0)
|
|
{
|
|
// the write was not successful so return that 0 bytes were written
|
|
return 0;
|
|
}
|
|
num_left -= status;
|
|
s += status;
|
|
}
|
|
return num;
|
|
}
|
|
}
|
|
}
|
|
|
|
// input functions
|
|
int_type underflow(
|
|
)
|
|
{
|
|
if (gptr() < egptr())
|
|
{
|
|
return static_cast<unsigned char>(*gptr());
|
|
}
|
|
|
|
int num_put_back = static_cast<int>(gptr() - eback());
|
|
if (num_put_back > max_putback)
|
|
{
|
|
num_put_back = max_putback;
|
|
}
|
|
|
|
// copy the putback characters into the putback end of the in_buffer
|
|
std::memmove(in_buffer+(max_putback-num_put_back), gptr()-num_put_back, num_put_back);
|
|
|
|
|
|
if (read_echoing_select(fd, fd_printf))
|
|
return EOF;
|
|
int num = read(fd,in_buffer+max_putback, in_buffer_size-max_putback);
|
|
if (num <= 0)
|
|
{
|
|
// an error occurred or the connection is over which is EOF
|
|
return EOF;
|
|
}
|
|
|
|
// reset in_buffer pointers
|
|
setg (in_buffer+(max_putback-num_put_back),
|
|
in_buffer+max_putback,
|
|
in_buffer+max_putback+num);
|
|
|
|
return static_cast<unsigned char>(*gptr());
|
|
}
|
|
|
|
std::streamsize xsgetn (
|
|
char_type* s,
|
|
std::streamsize n
|
|
)
|
|
{
|
|
std::streamsize temp = n;
|
|
while (n > 0)
|
|
{
|
|
int num = static_cast<int>(egptr() - gptr());
|
|
if (num >= n)
|
|
{
|
|
// copy data from our buffer
|
|
std::memcpy(s, gptr(), static_cast<size_t>(n));
|
|
gbump(static_cast<int>(n));
|
|
return temp;
|
|
}
|
|
|
|
// read more data into our buffer
|
|
if (num == 0)
|
|
{
|
|
if (underflow() == EOF)
|
|
break;
|
|
continue;
|
|
}
|
|
|
|
// copy all the data from our buffer
|
|
std::memcpy(s, gptr(), num);
|
|
n -= num;
|
|
gbump(num);
|
|
s += num;
|
|
}
|
|
return temp-n;
|
|
}
|
|
|
|
private:
|
|
|
|
// member data
|
|
int fd;
|
|
int fd_printf;
|
|
static const std::streamsize max_putback = 4;
|
|
static const std::streamsize out_buffer_size = 10000;
|
|
static const std::streamsize in_buffer_size = 10000;
|
|
char* out_buffer;
|
|
char* in_buffer;
|
|
|
|
};
|
|
|
|
namespace impl
|
|
{
|
|
int get_data_fd()
|
|
{
|
|
char* env_fd = getenv("DLIB_SUBPROCESS_DATA_FD");
|
|
DLIB_CASSERT(env_fd != 0,"");
|
|
return atoi(env_fd);
|
|
}
|
|
|
|
std::iostream& get_data_iostream()
|
|
{
|
|
static filestreambuf dbuff(get_data_fd(), -1);
|
|
static iostream out(&dbuff);
|
|
return out;
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
subprocess_stream::
|
|
subprocess_stream(const char* program_name) : stderr(NULL), iosub(NULL)
|
|
{
|
|
if (access(program_name, F_OK))
|
|
throw dlib::error("Error: '" + std::string(program_name) + "' file does not exist.");
|
|
if (access(program_name, X_OK))
|
|
throw dlib::error("Error: '" + std::string(program_name) + "' file is not executable.");
|
|
|
|
child_pid = fork();
|
|
if (child_pid == -1)
|
|
throw dlib::error("Failed to start child process");
|
|
|
|
if (child_pid == 0)
|
|
{
|
|
// In child process
|
|
dup2(stdout_pipe.child_fd(), STDOUT_FILENO);
|
|
dup2(stderr_pipe.child_fd(), STDERR_FILENO);
|
|
stdout_pipe.close();
|
|
stderr_pipe.close();
|
|
|
|
char* argv[] = {(char*)program_name, nullptr};
|
|
char* cudadevs = getenv("CUDA_VISIBLE_DEVICES");
|
|
if (cudadevs)
|
|
{
|
|
std::ostringstream sout;
|
|
sout << "DLIB_SUBPROCESS_DATA_FD="<<data_pipe.child_fd();
|
|
std::string extra = sout.str();
|
|
|
|
std::string extra2 = std::string("CUDA_VISIBLE_DEVICES=") + cudadevs;
|
|
char* envp[] = {(char*)extra.c_str(), (char*)extra2.c_str(), nullptr};
|
|
execve(argv[0], argv, envp);
|
|
}
|
|
else
|
|
{
|
|
std::ostringstream sout;
|
|
sout << "DLIB_SUBPROCESS_DATA_FD="<<data_pipe.child_fd();
|
|
std::string extra = sout.str();
|
|
char* envp[] = {(char*)extra.c_str(), nullptr};
|
|
execve(argv[0], argv, envp);
|
|
}
|
|
|
|
|
|
// If launching the child didn't work then bail immediately so the parent
|
|
// process has no chance to get tweaked out (*cough* MATLAB *cough*).
|
|
_Exit(1);
|
|
}
|
|
else
|
|
{
|
|
// In parent process
|
|
close(data_pipe.child_fd());
|
|
close(stdout_pipe.child_fd());
|
|
close(stderr_pipe.child_fd());
|
|
make_fd_non_blocking(data_pipe.parent_fd());
|
|
make_fd_non_blocking(stdout_pipe.parent_fd());
|
|
make_fd_non_blocking(stderr_pipe.parent_fd());
|
|
inout_buf = std::unique_ptr<filestreambuf>(new filestreambuf(data_pipe.parent_fd(), stdout_pipe.parent_fd()));
|
|
err_buf = std::unique_ptr<filestreambuf>(new filestreambuf(stderr_pipe.parent_fd(), stdout_pipe.parent_fd()));
|
|
iosub.rdbuf(inout_buf.get());
|
|
stderr.rdbuf(err_buf.get());
|
|
iosub.tie(&iosub);
|
|
stderr.tie(&iosub);
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
subprocess_stream::
|
|
~subprocess_stream()
|
|
{
|
|
try
|
|
{
|
|
wait();
|
|
}
|
|
catch (dlib::error& e)
|
|
{
|
|
std::cerr << e.what() << std::endl;
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
void subprocess_stream::
|
|
wait()
|
|
{
|
|
if (!wait_called)
|
|
{
|
|
wait_called = true;
|
|
send_eof();
|
|
|
|
std::ostringstream sout;
|
|
sout << stderr.rdbuf();
|
|
|
|
try{check_for_matlab_ctrl_c();} catch(...)
|
|
{
|
|
kill(child_pid, SIGTERM);
|
|
}
|
|
|
|
int status;
|
|
waitpid(child_pid, &status, 0);
|
|
if (status)
|
|
throw dlib::error("Child process terminated with an error.\n" + sout.str());
|
|
|
|
if (sout.str().size() != 0)
|
|
throw dlib::error("Child process terminated with an error.\n" + sout.str());
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
void subprocess_stream::
|
|
send_eof() { inout_buf->sync(); ::close(data_pipe.parent_fd()); }
|
|
|
|
// ----------------------------------------------------------------------------------------
|
|
|
|
}
|
|
|
|
|