From 87d2da676a458465eff55c67ae3a79c1eec685ab Mon Sep 17 00:00:00 2001
From: Craig Ringer <craig.ringer@enterprisedb.com>
Date: Thu, 13 Apr 2023 16:24:06 +1200
Subject: [PATCH] plugins: exec: exit fluent-bit after oneshot

Exit fluent-bit when the child process of the 'exec' plugin exits in
one-shot mode if the new Exit_After_Oneshot plugin option is true.

While exiting after the oneshot command is the logical behaviour, doing
so by default would change existing behaviour that may affect existing
uses, so it's not made the default.

Optionally propagate the child process exit code to the exit code of
fluent-bit itself if the new plugin option Propagate_Oneshot_Exit_Code
is True. This allows fluent-bit to be invoked as a wrapper for a child
command.

Signed-off-by: Craig Ringer <craig.ringer@enterprisedb.com>
---
 plugins/in_exec/in_exec.c | 62 ++++++++++++++++++++++++++++++++++++++-
 plugins/in_exec/in_exec.h |  2 ++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/plugins/in_exec/in_exec.c b/plugins/in_exec/in_exec.c
index 875b6139037..1d3597a593c 100644
--- a/plugins/in_exec/in_exec.c
+++ b/plugins/in_exec/in_exec.c
@@ -29,6 +29,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/wait.h>
 
 #include "in_exec.h"
 
@@ -37,6 +38,8 @@ static int in_exec_collect(struct flb_input_instance *ins,
                            struct flb_config *config, void *in_context)
 {
     int ret = -1;
+    int cmdret;
+    int decoded_cmdret;
     uint64_t val;
     size_t str_len = 0;
     FILE *cmdp = NULL;
@@ -165,7 +168,39 @@ static int in_exec_collect(struct flb_input_instance *ins,
 
  collect_end:
     if(cmdp != NULL){
-        pclose(cmdp);
+        cmdret = pclose(cmdp);
+        decoded_cmdret = 0;
+        if (cmdret == -1) {
+            flb_errno();
+        } else if (WIFEXITED(cmdret)) {
+            flb_plg_debug(ctx->ins, "command %s exited with code %d", ctx->cmd, WEXITSTATUS(cmdret));
+            decoded_cmdret = WEXITSTATUS(cmdret);
+        } else if (WIFSIGNALED(cmdret)) {
+            flb_plg_debug(ctx->ins, "command %s exited with signal %d", ctx->cmd, WTERMSIG(cmdret));
+            // There's no easy way to propagate a signal exit, so use a
+            // negative return code.
+            decoded_cmdret = -WTERMSIG(cmdret);
+        } else {
+            flb_plg_debug(ctx->ins, "command %s exited with unknown status", ctx->cmd);
+        }
+
+        /*
+         * In one-shot mode, exit fluent-bit once the child process terminates.
+         */
+        if (ctx->exit_after_oneshot == FLB_TRUE) {
+            /*
+             * propagate the child process exit code as the fluent-bit exit
+             * code so fluent-bit with the exec plugin can be used as a
+             * command wrapper.
+             */
+            if (ctx->propagate_exit_code == FLB_TRUE) {
+                config->exit_status_code = decoded_cmdret;
+            }
+            flb_plg_debug(ctx->ins, "one-shot command exited, terminating fluent-bit");
+            flb_engine_exit(config);
+        } else {
+            flb_plg_debug(ctx->ins, "one-shot command exited but exit_after_oneshot not set");
+        }
     }
 
     return ret;
@@ -212,6 +247,21 @@ static int in_exec_config_read(struct flb_exec *ctx,
         ctx->interval_nsec = atoi(DEFAULT_INTERVAL_NSEC);
     }
 
+    /*
+     * propagate_exit_code is not being forced to imply exit_after_oneshot in
+     * case somebody in future wishes to make the exec plugin exit on nonzero
+     * exit codes for normal repeating commands.
+     */
+    if (ctx->propagate_exit_code && !ctx->exit_after_oneshot) {
+        flb_plg_error(in, "propagate_exit_code=True option makes no sense without exit_after_oneshot=True");
+        return -1;
+    }
+
+    if (ctx->exit_after_oneshot && !ctx->oneshot) {
+        flb_plg_debug(in, "exit_after_oneshot implies oneshot mode, enabling");
+        ctx->oneshot = FLB_TRUE;
+    }
+
     if (ctx->oneshot) {
         ctx->interval_sec = -1;
         ctx->interval_nsec = -1;
@@ -378,6 +428,16 @@ static struct flb_config_map config_map[] = {
       0, FLB_TRUE, offsetof(struct flb_exec, oneshot),
       "execute the command only once"
     },
+    {
+      FLB_CONFIG_MAP_BOOL, "exit_after_oneshot", "false",
+      0, FLB_TRUE, offsetof(struct flb_exec, exit_after_oneshot),
+      "exit fluent-bit after the command terminates in one-shot mode"
+    },
+    {
+      FLB_CONFIG_MAP_BOOL, "propagate_exit_code", "false",
+      0, FLB_TRUE, offsetof(struct flb_exec, propagate_exit_code),
+      "propagate oneshot exit command fluent-bit exit code"
+    },
     /* EOF */
     {0}
 };
diff --git a/plugins/in_exec/in_exec.h b/plugins/in_exec/in_exec.h
index 64430b89be7..efde8c8d9ce 100644
--- a/plugins/in_exec/in_exec.h
+++ b/plugins/in_exec/in_exec.h
@@ -45,6 +45,8 @@ struct flb_exec {
     int interval_sec;
     int interval_nsec;
     struct flb_log_event_encoder log_encoder;
+    int exit_after_oneshot;
+    int propagate_exit_code;
 };
 
 #endif /* FLB_IN_EXEC_H */