[Project] Lua_config: Add post init scripts

1 files changed, 36 insertions, 7 deletions

diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index cff27ca1e..832e10614 100644
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -553,6 +553,7 @@ LUA_FUNCTION_DEF (config, register_worker_script);
 /***
  * @method rspamd_config:add_on_load(function(cfg, ev_base, worker) ... end)
  * Registers the following script to be executed when configuration is completely loaded
+ * and the worker is already started (forked)
  * @param {function} script function to be executed
  * @example
 rspamd_config:add_on_load(function(cfg, ev_base, worker)
@@ -584,6 +585,14 @@ end)
 LUA_FUNCTION_DEF (config, add_periodic);
 
 /***
+ * @method rspamd_config:add_post_init(function(cfg) ... end)
+ * Registers the following script to be executed when configuration is completely loaded
+ * @available 2.0+
+ * @param {function} script function to be executed
+ */
+LUA_FUNCTION_DEF (config, add_post_init);
+
+/***
  * @method rspamd_config:get_symbols_count()
  * Returns number of symbols registered in rspamd configuration
  * @return {number} number of symbols registered in the configuration
@@ -833,6 +842,7 @@ static const struct luaL_reg configlib_m[] = {
 	LUA_INTERFACE_DEF (config, register_re_selector),
 	LUA_INTERFACE_DEF (config, add_on_load),
 	LUA_INTERFACE_DEF (config, add_periodic),
+	LUA_INTERFACE_DEF (config, add_post_init),
 	LUA_INTERFACE_DEF (config, get_symbols_count),
 	LUA_INTERFACE_DEF (config, get_symbols_cksum),
 	LUA_INTERFACE_DEF (config, get_symbols_counters),
@@ -2894,16 +2904,35 @@ lua_config_add_on_load (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_config *cfg = lua_check_config (L, 1);
-	struct rspamd_config_post_load_script *sc;
+	struct rspamd_config_cfg_lua_script *sc;
+
+	if (cfg == NULL || lua_type (L, 2) != LUA_TFUNCTION) {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	sc = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (*sc));
+	lua_pushvalue (L, 2);
+	sc->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+	DL_APPEND (cfg->on_load_scripts, sc);
+
+	return 0;
+}
+
+static gint
+lua_config_add_post_init (lua_State *L)
+{
+	LUA_TRACE_POINT;
+	struct rspamd_config *cfg = lua_check_config (L, 1);
+	struct rspamd_config_cfg_lua_script *sc;
 
 	if (cfg == NULL || lua_type (L, 2) != LUA_TFUNCTION) {
 		return luaL_error (L, "invalid arguments");
 	}
 
-	sc = g_malloc0 (sizeof (*sc));
+	sc = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (*sc));
 	lua_pushvalue (L, 2);
 	sc->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
-	DL_APPEND (cfg->on_load, sc);
+	DL_APPEND (cfg->post_init_scripts, sc);
 
 	return 0;
 }
@@ -3315,13 +3344,13 @@ lua_config_register_finish_script (lua_State *L)
 {
 	LUA_TRACE_POINT;
 	struct rspamd_config *cfg = lua_check_config (L, 1);
-	struct rspamd_config_post_load_script *sc;
+	struct rspamd_config_cfg_lua_script *sc;
 
 	if (cfg != NULL && lua_type (L, 2) == LUA_TFUNCTION) {
-		sc = g_malloc0 (sizeof (*sc));
+		sc = rspamd_mempool_alloc0 (cfg->cfg_pool, sizeof (*sc));
 		lua_pushvalue (L, 2);
 		sc->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
-		DL_APPEND (cfg->finish_callbacks, sc);
+		DL_APPEND (cfg->on_term_scripts, sc);
 	}
 	else {
 		return luaL_error (L, "invalid arguments");
@@ -4033,7 +4062,7 @@ luaopen_config (lua_State * L)
 }
 
 void
-lua_call_finish_script (struct rspamd_config_post_load_script *sc,
+lua_call_finish_script (struct rspamd_config_cfg_lua_script *sc,
 		struct rspamd_task *task) {
 
 	struct rspamd_task **ptask;