static int output_handler (zio_t *z, const char *json_str, int len, void *arg)
{
struct subprocess *p = (struct subprocess *) arg;
json_object *o;
if (p->io_cb) {
if (!(o = json_tokener_parse (json_str))) {
errno = EINVAL;
return -1;
}
Jadd_int (o, "pid", subprocess_pid (p));
Jadd_str (o, "type", "io");
Jadd_str (o, "name", zio_name (z));
p->io_cb (p, json_object_to_json_string (o));
json_object_put (o);
}
else
send_output_to_stream (zio_name (z), json_str);
/*
* Check for process completion in case EOF from I/O stream and process
* already registered exit
*/
check_completion (p);
return (0);
}
static int forkzio_close_cb (zio_t *zio, void *arg)
{
forkzio_t ctx = arg;
ctx->readers--;
if ((ctx->flags & FORKZIO_FLAG_DEBUG))
msg ("%s: closing %s, %d readers left", __FUNCTION__,
zio_name (zio), ctx->readers);
return (ctx->readers > 0 ? 0 : -1); /* exit zloop when readers == 0 */
}
/* Data is ready on the zmq pair socket.
* Look for a zio matching the stream name and send it.
*/
static int forkzio_zsock_cb (zloop_t *zl, zmq_pollitem_t *zp, void *arg)
{
forkzio_t ctx = arg;
char *json_str = NULL;
zmsg_t *zmsg;
int i, rc = -1;
char *stream = NULL;
if (!(zmsg = zmsg_recv (zp->socket)))
goto done;
if (!(stream = zmsg_popstr (zmsg)))
goto done;
if (!(json_str = zmsg_popstr (zmsg)))
goto done;
for (i = 0; i < 3; i++) {
if (!ctx->zio[i] || strcmp (zio_name (ctx->zio[i]), stream) != 0)
continue;
if ((ctx->flags & FORKZIO_FLAG_DEBUG))
msg ("%s: msg %s => zio[%d]", __FUNCTION__, stream, i);
if (zio_write_json (ctx->zio[i], json_str) < 0) {
err ("zio_write_json");
goto done;
}
break;
}
/* N.B. if we wrote json containing only the eof and no data,
* our close callback will be called but from zio_write_json(),
* not from the reactor, so a -1 return from it won't cause the
* reactor to exit. Therefore, we have to catch the termination
* condition here (no more readers) and return -1 to the zloop.
*/
rc = ctx->readers > 0 ? 0 : -1;
done:
zmsg_destroy (&zmsg);
if (json_str)
free (json_str);
if (stream)
free (stream);
return rc;
}
int main (int argc, char **argv)
{
zio_t *zio;
int init_fds;
const char *name;
struct counts c;
int fd;
flux_reactor_t *r;
flux_watcher_t *w;
memset (&c, 0, sizeof (c));
plan (NO_PLAN);
test_encode ();
ok ((r = flux_reactor_create (0)) != NULL,
"flux reactor created");
init_fds = fdcount ();
diag ("initial fd count: %d", init_fds);
/* simple reader tests
*/
ok ((zio = zio_pipe_reader_create ("test1", &c)) != NULL,
"reader: zio_pipe_reader_create works");
ok ((name = zio_name (zio)) != NULL && !strcmp (name, "test1"),
"reader: zio_name returns correct name");
ok (zio_set_close_cb (zio, close_reader) == 0,
"reader: zio_set_close_cb works");
ok (zio_set_send_cb (zio, send_reader) == 0,
"reader: zio_set_send_cb works");
ok (zio_reactor_attach (zio, r) == 0,
"reader: zio_reactor_attach works");
ok ((fd = zio_dst_fd (zio)) >= 0,
"reader: zio_dst_fd returned valid file descriptor");
ok (write (fd, "narf!", 5) == 5,
"reader: wrote narf! to reader pipe");
ok (zio_close_dst_fd (zio) == 0,
"reader: zio_close_dst_fd succeeded");
ok (flux_reactor_run (r, 0) == 0,
"reader: reactor completed successfully");
ok (c.send_reader == 1,
"reader: send function called once for EOF + incomplete line");
errno = 0;
zio_destroy (zio);
ok (init_fds == fdcount (),
"reader: zio_destroy leaks no file descriptors");
/* simple writer tests
*/
ok ((zio = zio_pipe_writer_create ("test2", &c)) != NULL,
"writer: zio_pipe_writer_create works");
ok ((name = zio_name (zio)) != NULL && !strcmp (name, "test2"),
"writer: zio_name returns correct name");
ok (zio_set_close_cb (zio, close_writer) == 0,
"writer: zio_set_close_cb works");
ok ((fd = zio_src_fd (zio)) >= 0,
"writer: zio_src_fd returned valid file descriptor");
w = flux_fd_watcher_create (r, fd, FLUX_POLLIN, fd_read, &c);
ok (w != NULL,
"writer: created fd watcher");
flux_watcher_start (w);
ok (zio_write (zio, "narf!", 5) == 5,
"writer: zio_write narf! works");
ok (zio_write_eof (zio) == 0,
"writer: zio_write_eof works");
ok (flux_reactor_run (r, 0) == 0,
"writer: reactor completed successfully");
ok (c.fd_read_errors == 0 && c.fd_read_data == 5 && c.fd_read_eof == 1,
"writer: read narf + EOF on read end of pipe");
ok (c.close_writer == 1,
"writer: close callback invoked");
zio_destroy (zio);
ok (init_fds == fdcount (),
"writer: zio_destroy leaks no file descriptors");
flux_watcher_destroy (w);
flux_reactor_destroy (r);
done_testing ();
}
static void forkzio_pipe_thd (void *args, zctx_t *zctx, void *zs)
{
forkzio_t ctx = args;
zmq_pollitem_t zp = { .fd = -1, .socket = zs, .events = ZMQ_POLLIN };
zloop_t *zloop;
pid_t pid;
if (!(zloop = zloop_new ()))
oom ();
/* child stdin <= zs
*/
zloop_poller (zloop, &zp, (zloop_fn *)forkzio_zsock_cb, ctx);
ctx->zio[0] = zio_pipe_writer_create ("stdin", NULL);
if (zio_zloop_attach (ctx->zio[0], zloop) < 0)
err_exit ("zio_zloop_attach %s", zio_name (ctx->zio[0]));
/* child stdout => zs
*/
ctx->zio[1] = zio_pipe_reader_create ("stdout", zs, ctx);
zio_set_close_cb (ctx->zio[1], forkzio_close_cb);
if (zio_zloop_attach (ctx->zio[1], zloop) < 0)
err_exit ("zio_zloop_attach %s", zio_name (ctx->zio[1]));
ctx->readers++;
/* child stderr => zs
*/
ctx->zio[2] = zio_pipe_reader_create ("stderr", zs, ctx);
zio_set_close_cb (ctx->zio[2], forkzio_close_cb);
if (zio_zloop_attach (ctx->zio[2], zloop) < 0)
err_exit ("zio_zloop_attach %s", zio_name (ctx->zio[2]));
ctx->readers++;
pid = forkzio_fork (ctx);
(void)zloop_start (zloop);
forkzio_wait (pid);
zio_destroy (ctx->zio[0]);
zio_destroy (ctx->zio[1]);
zio_destroy (ctx->zio[2]);
zstr_send (zs, ""); /* signify EOF by sending an empty message */
}
static void forkzio_pty_thd (void *args, zctx_t *zctx, void *zs)
{
forkzio_t ctx = args;
zmq_pollitem_t zp = { .fd = -1, .socket = zs, .events = ZMQ_POLLIN };
zloop_t *zloop;
pid_t pid;
int ptyfd;
if (!(zloop = zloop_new ()))
oom ();
switch ((pid = forkpty (&ptyfd, NULL, NULL, NULL))) {
case -1: /* error */
err_exit ("forkpty");
case 0: /* child */
(void)execvp (ctx->av[0], ctx->av);
err_exit ("%s", ctx->av[0]);
default: /* parent */
break;
}
/* Data read from zs is written to pty master
*/
zloop_poller (zloop, &zp, (zloop_fn *)forkzio_zsock_cb, ctx);
ctx->zio[0] = zio_writer_create ("stdin", ptyfd, NULL);
zio_set_unbuffered (ctx->zio[0]);
if (zio_zloop_attach (ctx->zio[0], zloop) < 0)
err_exit ("zio_zloop_attach %s", zio_name (ctx->zio[0]));
/* Data read from pty master is written to zs
*/
ctx->zio[1] = zio_reader_create ("stdout", ptyfd, zs, ctx);
zio_set_unbuffered (ctx->zio[1]);
zio_set_close_cb (ctx->zio[1], forkzio_close_cb);
if (zio_zloop_attach (ctx->zio[1], zloop) < 0)
err_exit ("zio_zloop_attach %s", zio_name (ctx->zio[1]));
ctx->readers++;
(void)zloop_start (zloop);
forkzio_wait (pid);
zio_destroy (ctx->zio[0]);
zio_destroy (ctx->zio[1]);
zstr_send (zs, ""); /* signify EOF by sending an empty message */
}
forkzio_t forkzio_open (zctx_t *zctx, int ac, char **av, int flags)
{
zthread_attached_fn *thd = forkzio_pipe_thd;
forkzio_t ctx = xzmalloc (sizeof (*ctx));
ctx->ac = ac;
ctx->av = av;
ctx->zctx = zctx;
ctx->flags = flags;
if ((ctx->flags & FORKZIO_FLAG_PTY))
thd = forkzio_pty_thd;
if (!(ctx->zs = zthread_fork (zctx, thd, ctx))) {
free (ctx);
ctx = NULL;
}
return ctx;
}