shell: pty: prevent eof before all pty data is captured

Problem: When the pty.capture shell option is used (either implicitly
or explicitly), an eof for stdout may be written to the output eventlog
before all pty data is read. This is particularly true when used with
pty.interactive, since in this mode the normal stdout file descriptor
for the task is closed immediately when the task is attached to
the pty.

When capturing pty output for the output eventlog, take a reference
on the task's stdout channel so that it doesn't close until the pty
has completed reading its data.

Author: grondo

src/shell/pty.c

@@ -71,9 +71,15 @@ static void pty_monitor (struct flux_pty *pty, void *data, int len)
     flux_plugin_arg_t *args;
     int rank;
 
-    /* Don't bother sending 0 length output */
-    if (len == 0)
+    /*  len == 0 indicates pty is closed. If there's a reference on
+     *  stdout, release it here
+     */
+    if (len == 0) {
+        flux_subprocess_t *p;
+        if ((p = flux_pty_aux_get (pty, "subprocess")))
+            flux_subprocess_channel_decref (p, "stdout");
         return;
+    }
 
     rank = ptr2int (flux_pty_aux_get (pty, "rank"));
     if (!(args = flux_plugin_arg_create ())
@@ -351,7 +357,11 @@ static int pty_init (flux_plugin_t *p,
          */
         if (capture || rank != 0) {
             if (flux_pty_aux_set (pty, "shell", shell, NULL) < 0
-                || flux_pty_aux_set (pty, "rank", int2ptr (rank), NULL) < 0) {
+                || flux_pty_aux_set (pty, "rank", int2ptr (rank), NULL) < 0
+                || flux_pty_aux_set (pty,
+                                     "capture",
+                                     int2ptr (capture),
+                                     NULL) < 0) {
                 shell_log_errno ("flux_pty_aux_set");
                 goto error;
             }
@@ -408,6 +418,38 @@ static int pty_task_exec (flux_plugin_t *p,
     return (0);
 }
 
+static int pty_task_fork (flux_plugin_t *p,
+                          const char *topic,
+                          flux_plugin_arg_t *args,
+                          void *arg)
+{
+    flux_shell_t *shell = flux_plugin_get_shell (p);
+    flux_shell_task_t *task;
+    struct flux_pty *pty;
+    int rank;
+
+    if (!shell)
+        return shell_log_errno ("failed to get shell object");
+
+    if (flux_shell_getopt (shell, "pty", NULL) != 1)
+        return 0;
+
+    if (!(task = flux_shell_current_task (shell))
+        || flux_shell_task_info_unpack (task, "{s:i}", "rank", &rank) < 0)
+        return shell_log_errno ("unable to get task rank");
+
+    /*  If pty is in capture mode, then take a reference on subprocess
+     *  stdout so that EOF is not read until pty exits.
+     */
+   if ((pty = pty_lookup (shell, rank))
+       && ptr2int (flux_pty_aux_get (pty, "capture"))) {
+       flux_subprocess_t *sp = flux_shell_task_subprocess (task);
+       flux_subprocess_channel_incref (sp, "stdout");
+       flux_pty_aux_set (pty, "subprocess", sp, NULL);
+   }
+   return (0);
+}
+
 static int pty_task_exit (flux_plugin_t *p,
                           const char *topic,
                           flux_plugin_arg_t *args,
@@ -446,6 +488,7 @@ struct shell_builtin builtin_pty = {
     .name = FLUX_SHELL_PLUGIN_NAME,
     .init = pty_init,
     .task_exec = pty_task_exec,
+    .task_fork = pty_task_fork,
     .task_exit = pty_task_exit,
 };