From bb9ff31dc7fb4378fde731388a32b9de1d698ebe Mon Sep 17 00:00:00 2001
From: Hiroshi Hatake <hatake@calyptia.com>
Date: Wed, 27 Sep 2023 18:32:49 +0900
Subject: [PATCH] plugin_proxy: Add pre registration callback for hot-reloading
 state handling

Signed-off-by: Hiroshi Hatake <hatake@calyptia.com>
---
 include/fluent-bit/flb_config.h |  1 +
 src/flb_config.c                |  1 +
 src/flb_plugin_proxy.c          | 11 +++++++++++
 src/flb_reload.c                |  3 +++
 4 files changed, 16 insertions(+)

diff --git a/include/fluent-bit/flb_config.h b/include/fluent-bit/flb_config.h
index 85e5b77984f..f2392901db8 100644
--- a/include/fluent-bit/flb_config.h
+++ b/include/fluent-bit/flb_config.h
@@ -262,6 +262,7 @@ struct flb_config {
     int ensure_thread_safety_on_hot_reloading;
     unsigned int hot_reloaded_count;
     int shutdown_by_hot_reloading;
+    int hot_reloading;
 
     /* Co-routines */
     unsigned int coro_stack_size;
diff --git a/src/flb_config.c b/src/flb_config.c
index aa65a9eb266..ae19fcbf7ec 100644
--- a/src/flb_config.c
+++ b/src/flb_config.c
@@ -282,6 +282,7 @@ struct flb_config *flb_config_init()
     config->ensure_thread_safety_on_hot_reloading = FLB_TRUE;
     config->hot_reloaded_count = 0;
     config->shutdown_by_hot_reloading = FLB_FALSE;
+    config->hot_reloading = FLB_FALSE;
 
 #ifdef FLB_HAVE_SQLDB
     mk_list_init(&config->sqldb_list);
diff --git a/src/flb_plugin_proxy.c b/src/flb_plugin_proxy.c
index e8be283b52d..eea90729762 100644
--- a/src/flb_plugin_proxy.c
+++ b/src/flb_plugin_proxy.c
@@ -424,8 +424,19 @@ int flb_plugin_proxy_register(struct flb_plugin_proxy *proxy,
 {
     int ret;
     int (*cb_register)(struct flb_plugin_proxy_def *);
+    int (*cb_pre_register)(int);
     struct flb_plugin_proxy_def *def = proxy->def;
 
+    /* Lookup the pre registration callback */
+    cb_pre_register = flb_plugin_proxy_symbol(proxy, "FLBPluginPreRegister");
+    if (cb_pre_register != NULL) {
+        /* Prepare the registration if available */
+        ret = cb_pre_register(config->hot_reloading);
+        if (ret == -1) {
+            return -1;
+        }
+    }
+
     /* Lookup the registration callback */
     cb_register = flb_plugin_proxy_symbol(proxy, "FLBPluginRegister");
     if (!cb_register) {
diff --git a/src/flb_reload.c b/src/flb_reload.c
index 7b113d78204..6c0e7c6ce90 100644
--- a/src/flb_reload.c
+++ b/src/flb_reload.c
@@ -430,6 +430,8 @@ int flb_reload(flb_ctx_t *ctx, struct flb_cf *cf_opts)
     reloaded_count = ctx->config->hot_reloaded_count + 1;
     /* Mark shutdown reason as hot_reloading */
     ctx->config->shutdown_by_hot_reloading = FLB_TRUE;
+    /* Mark hot reloading */
+    new_config->hot_reloading = FLB_TRUE;
 
 #ifdef FLB_HAVE_STREAM_PROCESSOR
     /* Inherit stream processor definitions from command line */
@@ -513,6 +515,7 @@ int flb_reload(flb_ctx_t *ctx, struct flb_cf *cf_opts)
     if (ret == 0) {
         new_config->hot_reloaded_count = reloaded_count;
         flb_debug("[reload] hot reloaded %d time(s)", reloaded_count);
+        new_config->hot_reloading = FLB_FALSE;
     }
 
     return 0;