tee: enhance -p mode using iopoll() to detect broken pipe outputs

edquist · pixelb · commit 6b12e62d9585 · 2023-02-28T14:02:42.000Z
If input is intermittent (a tty, pipe, or socket), and all remaining
outputs are pipes (eg, &gt;(cmd) process substitutions), exit early when
they have all become broken pipes (and thus future writes will fail),
without waiting for more input to become available, as future write
attempts to these outputs will fail (SIGPIPE/EPIPE).

Only provide this enhancement when pipe errors are ignored (-p mode).

Note that only one output needs to be monitored at a time with iopoll(),
as we only want to exit early if _all_ outputs have been removed.

* src/tee.c (pipe_check): New global for iopoll mode.
(main): enable pipe_check for -p, as long as output_error ignores EPIPE,
and input is suitable for iopoll().
(get_next_out): Helper function for finding next valid output.
(fail_output, tee_files): Break out write failure/output removal logic
to helper function.
(tee_files): Add out_pollable array to track which outputs are suitable
for iopoll() (ie, that are pipes); track first output index that is
still valid; add iopoll() broken pipe detection before calling read(),
removing an output that becomes a broken pipe.
* src/local.mk (src_tee_SOURCES): include src/iopoll.c.
* NEWS: Mention tee -p enhancement in Improvements.
* doc/coreutils.texi: Mention the new early exit behavior in the nopipe
modes for the tee -p option.

Suggested-by: Arsen Arsenović &lt;arsen@aarsen.me&gt;
diff --git a/NEWS b/NEWS
@@ -153,6 +153,10 @@ GNU coreutils NEWS                                    -*- outline -*-
   when their modification time doesn't change when new data is available.
   Previously tail would not show any new data in this case.
 
+  tee -p detects when all remaining outputs have become broken pipes, and
+  exits, rather than waiting for more input to induce an exit when written.
+
+
 
 * Noteworthy changes in release 9.1 (2022-04-15) [stable]
 
diff --git a/doc/coreutils.texi b/doc/coreutils.texi
@@ -14198,13 +14198,15 @@ This is the default @var{mode} when not specified,
 or when the short form @option{-p} is used.
 Warn on error opening or writing any output, except pipes.
 Writing is continued to still open files/pipes.
+Exit early if all remaining outputs become broken pipes.
 Exit status indicates failure if any non pipe output had an error.
 
 @item exit
 Exit on error opening or writing any output, including pipes.
 
 @item exit-nopipe
 Exit on error opening or writing any output, except pipes.
+Exit early if all remaining outputs become broken pipes.
 @end table
 
 @end table
diff --git a/src/iopoll.h b/src/iopoll.h
@@ -1,6 +1,6 @@
 #define IOPOLL_BROKEN_OUTPUT -2
 #define IOPOLL_ERROR         -3
 
-int iopoll(int fdin, int fdout);
-bool iopoll_input_ok(int fdin);
-bool iopoll_output_ok(int fdout);
+int iopoll (int fdin, int fdout);
+bool iopoll_input_ok (int fdin);
+bool iopoll_output_ok (int fdout);
diff --git a/src/local.mk b/src/local.mk
@@ -396,6 +396,7 @@ src_arch_SOURCES = src/uname.c src/uname-arch.c
 src_cut_SOURCES = src/cut.c src/set-fields.c
 src_numfmt_SOURCES = src/numfmt.c src/set-fields.c
 
+src_tee_SOURCES = src/tee.c src/iopoll.c
 src_sum_SOURCES = src/sum.c src/sum.h src/digest.c
 src_sum_CPPFLAGS = -DHASH_ALGO_SUM=1 $(AM_CPPFLAGS)
 
diff --git a/src/tee.c b/src/tee.c
@@ -28,6 +28,7 @@
 #include "fadvise.h"
 #include "stdio--.h"
 #include "xbinary-io.h"
+#include "iopoll.h"
 
 /* The official name of this program (e.g., no 'g' prefix).  */
 #define PROGRAM_NAME "tee"
@@ -45,6 +46,9 @@ static bool append;
 /* If true, ignore interrupts. */
 static bool ignore_interrupts;
 
+/* If true, detect if next output becomes broken while waiting for input. */
+static bool pipe_check;
+
 enum output_error
   {
     output_error_sigpipe,      /* traditional behavior, sigpipe enabled.  */
@@ -149,6 +153,12 @@ main (int argc, char **argv)
                                       output_error_args, output_error_types);
           else
             output_error = output_error_warn_nopipe;
+
+          /* Detect and close a broken pipe output when ignoring EPIPE.  */
+          if (output_error == output_error_warn_nopipe
+              || output_error == output_error_exit_nopipe)
+            pipe_check = true;
+
           break;
 
         case_GETOPT_HELP_CHAR;
@@ -166,6 +176,10 @@ main (int argc, char **argv)
   if (output_error != output_error_sigpipe)
     signal (SIGPIPE, SIG_IGN);
 
+  /* No need to poll outputs if input is always ready for reading.  */
+  if (pipe_check && !iopoll_input_ok (STDIN_FILENO))
+    pipe_check = false;
+
   /* Do *not* warn if tee is given no file arguments.
      POSIX requires that it work when given no arguments.  */
 
@@ -176,6 +190,42 @@ main (int argc, char **argv)
   return ok ? EXIT_SUCCESS : EXIT_FAILURE;
 }
 
+
+/* Return the index of the first non-NULL descriptor after idx,
+   or -1 if all are NULL.  */
+
+static int
+get_next_out (FILE **descriptors, int nfiles, int idx)
+{
+  for (idx++; idx <= nfiles; idx++)
+    if (descriptors[idx])
+      return idx;
+  return -1;  /* no outputs remaining */
+}
+
+/* Remove descriptors[i] due to write failure or broken pipe.
+   Return true if this indicates a reportable error.  */
+
+static bool
+fail_output (FILE **descriptors, char **files, int i)
+{
+  int w_errno = errno;
+  bool fail = errno != EPIPE
+              || output_error == output_error_exit
+              || output_error == output_error_warn;
+  if (descriptors[i] == stdout)
+    clearerr (stdout); /* Avoid redundant close_stdout diagnostic.  */
+  if (fail)
+    {
+      error (output_error == output_error_exit
+             || output_error == output_error_exit_nopipe,
+             w_errno, "%s", quotef (files[i]));
+    }
+  descriptors[i] = NULL;
+  return fail;
+}
+
+
 /* Copy the standard input into each of the NFILES files in FILES
    and into the standard output.  As a side effect, modify FILES[-1].
    Return true if successful.  */
@@ -185,9 +235,11 @@ tee_files (int nfiles, char **files)
 {
   size_t n_outputs = 0;
   FILE **descriptors;
+  bool *out_pollable;
   char buffer[BUFSIZ];
   ssize_t bytes_read = 0;
   int i;
+  int first_out = 0;  /* idx of first non-NULL output in descriptors */
   bool ok = true;
   char const *mode_string =
     (O_BINARY
@@ -202,8 +254,12 @@ tee_files (int nfiles, char **files)
      In both arrays, entry 0 corresponds to standard output.  */
 
   descriptors = xnmalloc (nfiles + 1, sizeof *descriptors);
+  if (pipe_check)
+    out_pollable = xnmalloc (nfiles + 1, sizeof *out_pollable);
   files--;
   descriptors[0] = stdout;
+  if (pipe_check)
+    out_pollable[0] = iopoll_output_ok (fileno (descriptors[0]));
   files[0] = bad_cast (_("standard output"));
   setvbuf (stdout, NULL, _IONBF, 0);
   n_outputs++;
@@ -212,6 +268,8 @@ tee_files (int nfiles, char **files)
     {
       /* Do not treat "-" specially - as mandated by POSIX.  */
       descriptors[i] = fopen (files[i], mode_string);
+      if (pipe_check)
+        out_pollable[i] = iopoll_output_ok (fileno (descriptors[i]));
       if (descriptors[i] == NULL)
         {
           error (output_error == output_error_exit
@@ -228,6 +286,28 @@ tee_files (int nfiles, char **files)
 
   while (n_outputs)
     {
+      if (pipe_check && out_pollable[first_out])
+        {
+          /* Monitor for input, or errors on first valid output.  */
+          int err = iopoll (STDIN_FILENO, fileno (descriptors[first_out]));
+
+          /* Close the output if it became a broken pipe.  */
+          if (err == IOPOLL_BROKEN_OUTPUT)
+            {
+              errno = EPIPE;  /* behave like write produced EPIPE */
+              if (fail_output (descriptors, files, first_out))
+                ok = false;
+              n_outputs--;
+              first_out = get_next_out (descriptors, nfiles, first_out);
+              continue;
+            }
+          else if (err == IOPOLL_ERROR)
+            {
+              error (0, errno, _("iopoll error"));
+              ok = false;
+            }
+        }
+
       bytes_read = read (STDIN_FILENO, buffer, sizeof buffer);
       if (bytes_read < 0 && errno == EINTR)
         continue;
@@ -240,21 +320,11 @@ tee_files (int nfiles, char **files)
         if (descriptors[i]
             && fwrite (buffer, bytes_read, 1, descriptors[i]) != 1)
           {
-            int w_errno = errno;
-            bool fail = errno != EPIPE || (output_error == output_error_exit
-                                          || output_error == output_error_warn);
-            if (descriptors[i] == stdout)
-              clearerr (stdout); /* Avoid redundant close_stdout diagnostic.  */
-            if (fail)
-              {
-                error (output_error == output_error_exit
-                       || output_error == output_error_exit_nopipe,
-                       w_errno, "%s", quotef (files[i]));
-              }
-            descriptors[i] = NULL;
-            if (fail)
+            if (fail_output (descriptors, files, i))
               ok = false;
             n_outputs--;
+            if (i == first_out)
+              first_out = get_next_out (descriptors, nfiles, first_out);
           }
     }
 
@@ -273,6 +343,8 @@ tee_files (int nfiles, char **files)
       }
 
   free (descriptors);
+  if (pipe_check)
+    free (out_pollable);
 
   return ok;
 }
