Merge pull request #5155 from grondo/output-truncate

mergify[bot] · web-flow · commit 6d95773b8fcc · 2023-05-08T16:20:23.000Z
shell: truncate output to KVS after 10MB
diff --git a/src/common/libioencode/ioencode.c b/src/common/libioencode/ioencode.c
@@ -102,11 +102,16 @@ static int decode_data_base64 (char *src,
 {
     ssize_t rc;
     size_t size = base64_decoded_length (srclen);
-    if (!(*datap = malloc (size)))
-        return -1;
-    if ((rc = base64_decode (*datap, size, src, srclen)) < 0)
-        return -1;
-    *lenp = rc;
+    if (datap) {
+        if (!(*datap = malloc (size)))
+            return -1;
+        if ((rc = base64_decode (*datap, size, src, srclen)) < 0)
+            return -1;
+        if (lenp)
+            *lenp = rc;
+    }
+    else if (lenp)
+        *lenp = size;
     return 0;
 }
 
@@ -154,18 +159,21 @@ int iodecode (json_t *o,
         (*streamp) = stream;
     if (rankp)
         (*rankp) = rank;
-    if (datap) {
-        *datap = NULL;
+    if (datap || lenp) {
+        if (datap)
+            *datap = NULL;
         if (data) {
             if (encoding && strcmp (encoding, "base64") == 0) {
                 if (decode_data_base64 (data, len, datap, &bin_len) < 0)
                     goto cleanup;
             }
             else {
                 bin_len = len;
-                if (!(*datap = malloc (bin_len)))
-                    goto cleanup;
-                memcpy (*datap, data, bin_len);
+                if (datap) {
+                    if (!(*datap = malloc (bin_len)))
+                        goto cleanup;
+                    memcpy (*datap, data, bin_len);
+                }
             }
         }
     }
diff --git a/src/common/libioencode/ioencode.h b/src/common/libioencode/ioencode.h
@@ -29,6 +29,7 @@ json_t *ioencode (const char *stream,
  * - both data and EOF can be available
  * - if no data available, data set to NULL and len to 0
  * - data must be freed after return
+ * - data can be NULL and len non-NULL to retrieve data length
  */
 int iodecode (json_t *o,
               const char **stream,
diff --git a/src/common/libioencode/test/ioencode.c b/src/common/libioencode/test/ioencode.c
@@ -62,6 +62,14 @@ void basic (void)
         && eof == true,
         "iodecode returned correct info");
     free (data);
+
+    ok (iodecode (o, &stream, &rank, NULL, &len, &eof) == 0,
+        "iodecode can be passed NULL data to query len");
+    ok (!strcmp (stream, "stdout")
+        && !strcmp (rank, "[0-8]")
+        && len == 3
+        && eof == true,
+        "iodecode returned correct info");
     json_decref (o);
 
     ok ((o = ioencode ("stderr", "[4,5]", NULL, 0, true)) != NULL,
@@ -75,6 +83,14 @@ void basic (void)
         && eof == true,
         "iodecode returned correct info");
     free (data);
+
+    ok (iodecode (o, &stream, &rank, NULL, &len, &eof) == 0,
+        "iodecode can be passed NULL data to query len");
+    ok (!strcmp (stream, "stderr")
+        && !strcmp (rank, "[4,5]")
+        && len == 0
+        && eof == true,
+        "iodecode returned correct info");
     json_decref (o);
 }
 
@@ -108,6 +124,10 @@ static void binary_data (void)
     ok (memcmp (data, buffer, len) == 0,
         "data matches");
     free (data);
+    ok (iodecode (o, &stream, &rank, NULL, &len, &eof) == 0,
+        "iodecode can be passed NULL data to query len");
+    ok (len == sizeof (buffer),
+        "len is correct");
     json_decref (o);
 }
 
diff --git a/src/shell/output.c b/src/shell/output.c
@@ -62,6 +62,9 @@
 #include "builtins.h"
 #include "log.h"
 
+#define OUTPUT_LIMIT_BYTES 1024*1024*10
+#define OUTPUT_LIMIT_STRING "10MB"
+
 enum {
     FLUX_OUTPUT_TYPE_TERM = 1,
     FLUX_OUTPUT_TYPE_KVS = 2,
@@ -88,6 +91,8 @@ struct shell_output {
     bool stopped;
     int stdout_type;
     int stderr_type;
+    size_t stdout_bytes;
+    size_t stderr_bytes;
     struct shell_output_type_file stdout_file;
     struct shell_output_type_file stderr_file;
     zhash_t *fds;
@@ -298,36 +303,92 @@ static int shell_output_kvs_init (struct shell_output *out, json_t *header)
     return rc;
 }
 
-static int entry_output_is_kvs (struct shell_output *out, json_t *entry)
+/*  Return true if entry is a kvs destination, false otherwise.
+ *  If true, then then stream and len will be set to the stream and
+ *   length of data in this entry.
+ */
+static bool entry_output_is_kvs (struct shell_output *out,
+                                 json_t *entry,
+                                 bool *stdoutp,
+                                 int *lenp,
+                                 bool *eofp)
 {
     json_t *context;
     const char *name;
     const char *stream;
+
     if (eventlog_entry_parse (entry, NULL, &name, &context) < 0) {
         shell_log_errno ("eventlog_entry_parse");
         return 0;
     }
     if (!strcmp (name, "data")) {
-        if (iodecode (context, &stream, NULL, NULL, NULL, NULL) < 0) {
+        if (iodecode (context, &stream, NULL, NULL, lenp, eofp) < 0) {
             shell_log_errno ("iodecode");
             return 0;
         }
-        if (!strcmp (stream, "stdout"))
+        if ((*stdoutp = !strcmp (stream, "stdout")))
             return (out->stdout_type == FLUX_OUTPUT_TYPE_KVS);
         else
             return (out->stderr_type == FLUX_OUTPUT_TYPE_KVS);
     }
     return 0;
 }
 
+static bool check_kvs_output_limit (struct shell_output *out,
+                                    bool is_stdout,
+                                    int len)
+{
+    const char *stream;
+    size_t *bytesp;
+    size_t prev;
+
+    if (is_stdout) {
+        stream = "stdout";
+        bytesp = &out->stdout_bytes;
+    }
+    else {
+        stream = "stderr";
+        bytesp = &out->stderr_bytes;
+    }
+
+    prev = *bytesp;
+    *bytesp += len;
+
+    if (*bytesp > OUTPUT_LIMIT_BYTES) {
+        /*  Only log an error when the threshold is reached.
+        */
+        if (prev <= OUTPUT_LIMIT_BYTES)
+            shell_warn ("%s will be truncated, %s limit exceeded",
+                        stream,
+                        OUTPUT_LIMIT_STRING);
+        return true;
+    }
+    return false;
+}
+
 static int shell_output_kvs (struct shell_output *out)
 {
     json_t *entry;
     size_t index;
+    bool is_stdout;
+    int len;
+    bool eof;
+
     json_array_foreach (out->output, index, entry) {
-        if (entry_output_is_kvs (out, entry) &&
-            eventlogger_append_entry (out->ev, 0, "output", entry) < 0) {
-            return shell_log_errno ("eventlogger_append");
+        if (entry_output_is_kvs (out, entry, &is_stdout, &len, &eof)) {
+            bool truncate = check_kvs_output_limit (out, is_stdout, len);
+            if (!truncate || eof) {
+                if (eventlogger_append_entry (out->ev, 0, "output", entry) < 0)
+                    return shell_log_errno ("eventlogger_append");
+            }
+            if (eof && truncate) {
+                size_t total = is_stdout ?
+                               out->stdout_bytes : out->stderr_bytes;
+                shell_warn ("%s: %zu of %zu bytes truncated",
+                            is_stdout ? "stdout" : "stderr",
+                            total - OUTPUT_LIMIT_BYTES,
+                            total);
+            }
         }
     }
     return 0;
diff --git a/t/t2606-job-shell-output-redirection.t b/t/t2606-job-shell-output-redirection.t
@@ -348,4 +348,15 @@ test_expect_success 'job-shell: shell errors are captured in error file' '
 	test_expect_code 127  flux run --error=test.err nosuchcommand &&
 	grep "nosuchcommand: No such file or directory" test.err
 '
+test_expect_success LONGTEST 'job-shell: output to kvs is truncated at 10MB' '
+	dd if=/dev/urandom bs=10240 count=800 | base64 --wrap 79 >expected &&
+	flux run cat expected >output 2>truncate.error &&
+	test_debug "cat truncate.error" &&
+	grep "stdout.*truncated" truncate.error
+'
+test_expect_success LONGTEST 'job-shell: stderr to kvs is truncated at 10MB' '
+	dd if=/dev/urandom bs=10240 count=800 | base64 --wrap 79 >expected &&
+	flux run sh -c "cat expected >&2"  >truncate2.error 2>&1 &&
+	grep "stderr.*truncated" truncate2.error
+'
 test_done
