✨ feat(module manager): add thread_init and thread_exit API

6 files changed, 143 insertions, 50 deletions

diff --git a/include/stellar/module_manager.h b/include/stellar/module_manager.h
index 47e8384..8caf702 100644
--- a/include/stellar/module_manager.h
+++ b/include/stellar/module_manager.h
@@ -20,21 +20,25 @@ void stellar_module_set_name(struct stellar_module* mod, const char *name);
 
 struct stellar_module_manager;
 
-typedef struct stellar_module *module_on_init_func(struct stellar_module_manager *mod_mgr);
-typedef void module_on_exit_func(struct stellar_module_manager *mod_mgr, struct stellar_module *mod);
+typedef struct stellar_module *module_on_instance_init_func(struct stellar_module_manager *mod_mgr);
+typedef void module_on_instance_exit_func(struct stellar_module_manager *mod_mgr, struct stellar_module *mod);
+
+typedef struct stellar_module *module_on_thread_init_func(struct stellar_module_manager *mod_mgr, int thread_id, struct stellar_module *mod);
+typedef void module_on_thread_exit_func(struct stellar_module_manager *mod_mgr, int thread_id, struct stellar_module *mod);
 
 struct stellar_module_manager *stellar_module_manager_new(const char *module_spec_toml_path, int max_thread_num, struct mq_schema *mq_schema, struct logger *logger);
 void stellar_module_manager_free(struct stellar_module_manager *mod_mgr);
 
-void stellar_module_manager_register_thread(struct stellar_module_manager* mod_mgr, int thread_id, struct mq_runtime *mq_rt);
+void stellar_module_manager_register_thread(struct stellar_module_manager *mod_mgr, int thread_id, struct mq_runtime *mq_rt);
+void stellar_module_manager_unregister_thread(struct stellar_module_manager *mod_mgr, int thread_id);
 
 // return -1 on error
-int stellar_module_manager_get_thread_id(struct stellar_module_manager* mod_mgr);
+int stellar_module_manager_get_thread_id(struct stellar_module_manager *mod_mgr);
 struct mq_runtime *stellar_module_manager_get_mq_runtime(struct stellar_module_manager *mod_mgr);
 
 struct stellar_module *stellar_module_manager_get_module(struct stellar_module_manager *mod_mgr, const char *module_name); 
 
-int stellar_module_manager_get_max_thread_num(struct stellar_module_manager* mod_mgr);
+int stellar_module_manager_get_max_thread_num(struct stellar_module_manager *mod_mgr);
 const char *stellar_module_manager_get_toml_path(struct stellar_module_manager *mod_mgr);
 struct mq_schema *stellar_module_manager_get_mq_schema(struct stellar_module_manager *mod_mgr);
 struct logger *stellar_module_manager_get_logger(struct stellar_module_manager *mod_mgr);
diff --git a/infra/CMakeLists.txt b/infra/CMakeLists.txt
index d2bab28..be0f97b 100644
--- a/infra/CMakeLists.txt
+++ b/infra/CMakeLists.txt
@@ -9,13 +9,18 @@ foreach(infra ${INFRA})
 endforeach()
 
 add_library(stellar_lib SHARED stellar_core.c stellar_stat.c)
+add_dependencies(stellar_lib ${WHOLE_ARCHIVE})
+
 set_target_properties(stellar_lib PROPERTIES LINK_FLAGS "-Wl,--version-script=${CMAKE_CURRENT_LIST_DIR}/version.map")
 target_link_libraries(stellar_lib PRIVATE pthread -Wl,--whole-archive ${WHOLE_ARCHIVE} -Wl,--no-whole-archive)
+
 target_link_libraries(stellar_lib PUBLIC ${LIBS})
 target_link_options(stellar_lib PRIVATE -rdynamic)
 set_target_properties(stellar_lib PROPERTIES OUTPUT_NAME "stellar")
 
 add_executable(stellar main.c stellar_core.c stellar_stat.c)
+add_dependencies(stellar_lib ${WHOLE_ARCHIVE})
+
 set_target_properties(stellar PROPERTIES LINK_FLAGS "-Wl,--version-script=${CMAKE_CURRENT_LIST_DIR}/version.map")
 target_link_libraries(stellar PRIVATE pthread -Wl,--whole-archive ${WHOLE_ARCHIVE} -Wl,--no-whole-archive ${LIBS})
 target_link_options(stellar PRIVATE -rdynamic)
diff --git a/infra/module_manager/module_manager.c b/infra/module_manager/module_manager.c
index e5ac28c..b727ebb 100644
--- a/infra/module_manager/module_manager.c
+++ b/infra/module_manager/module_manager.c
@@ -29,9 +29,6 @@ struct stellar_module_manager *stellar_module_manager_new(const char *module_spe
     int mod_num = 0;
 	toml_table_t *conf = toml_parse_file(fp, NULL, 0);
     toml_array_t* mod_array = toml_array_in(conf, "module");
-    char *path = NULL;
-    char *init_func_name = NULL;
-    char *exit_func_name = NULL;
 
 	if(mod_array==NULL)goto MODULE_SPEC_LOAD_END;
 
@@ -45,44 +42,73 @@ struct stellar_module_manager *stellar_module_manager_new(const char *module_spe
         const char *path_raw = toml_raw_in(toml_mod, "path");
         const char *init_func_name_raw = toml_raw_in(toml_mod, "init");
         const char *exit_func_name_raw = toml_raw_in(toml_mod, "exit");
-        if (toml_rtos(path_raw, &path) || toml_rtos(init_func_name_raw, &init_func_name) ||
-            toml_rtos(exit_func_name_raw, &exit_func_name))
-        {
-            goto MODULE_SPEC_LOAD_END;
-        }
-        void* handle = dlopen(path, RTLD_NOW|RTLD_LAZY|RTLD_GLOBAL);
+		const char *thread_init_func_name_raw = toml_raw_in(toml_mod, "thread_init");
+		const char *thread_exit_func_name_raw = toml_raw_in(toml_mod, "thread_exit");
+
+		toml_rtos(path_raw, &mod_mgr->module_specs[i].path);
+		toml_rtos(init_func_name_raw, &mod_mgr->module_specs[i].instance_init_cb_name);
+		toml_rtos(exit_func_name_raw, &mod_mgr->module_specs[i].instance_exit_cb_name);
+		toml_rtos(thread_init_func_name_raw, &mod_mgr->module_specs[i].thread_init_cb_name);
+		toml_rtos(thread_exit_func_name_raw, &mod_mgr->module_specs[i].thread_exit_cb_name);
+
+        void* handle = dlopen(mod_mgr->module_specs[i].path, RTLD_NOW|RTLD_LAZY|RTLD_GLOBAL);
         if (!handle) {
-            fprintf(stderr, "Error loading module %s: %s\n", path, dlerror());
-            goto MODULE_SPEC_LOAD_END;
+            fprintf(stderr, "Error loading module %s: %s\n", mod_mgr->module_specs[i].path, dlerror());
+            break;
         }
-
-        mod_mgr->module_specs[i].on_init_cb = (module_on_init_func *) dlsym(handle, init_func_name);
-        if (mod_mgr->module_specs[i].on_init_cb)
+        if (mod_mgr->module_specs[i].instance_init_cb_name)
         {
-            mod_mgr->module_specs[i].mod = mod_mgr->module_specs[i].on_init_cb(mod_mgr);
-            if (stellar_module_manager_get_module(mod_mgr, mod_mgr->module_specs[i].mod->name) == NULL)
+            mod_mgr->module_specs[i].on_instance_init_cb =
+                (module_on_instance_init_func *)dlsym(handle, mod_mgr->module_specs[i].instance_init_cb_name);
+            if (mod_mgr->module_specs[i].on_instance_init_cb)
             {
-                mod_mgr->module_specs[i].is_init_succ = true;
+                mod_mgr->module_specs[i].mod = mod_mgr->module_specs[i].on_instance_init_cb(mod_mgr);
+                if (stellar_module_manager_get_module(mod_mgr, mod_mgr->module_specs[i].mod->name) == NULL)
+                {
+                    mod_mgr->module_specs[i].init_succ = true;
+                }
+                else
+                {
+                    fprintf(stderr, "Module %s already exists\n", mod_mgr->module_specs[i].mod->name);
+                    if (mod_mgr->module_specs[i].on_instance_exit_cb)
+                        mod_mgr->module_specs[i].on_instance_exit_cb(mod_mgr, mod_mgr->module_specs[i].mod);
+                }
             }
             else
             {
-                fprintf(stderr, "Module %s already exists\n", mod_mgr->module_specs[i].mod->name);
-                if (mod_mgr->module_specs[i].on_exit_cb)
-                    mod_mgr->module_specs[i].on_exit_cb(mod_mgr, mod_mgr->module_specs[i].mod);
+                fprintf(stderr, "Could not load init function %s: %s\n", mod_mgr->module_specs[i].instance_init_cb_name, dlerror());
             }
         }
-        else
+        if (mod_mgr->module_specs[i].instance_exit_cb_name)
         {
-            fprintf(stderr, "Could not load init function %s: %s\n", init_func_name, dlerror());
+            mod_mgr->module_specs[i].on_instance_exit_cb =
+                (module_on_instance_exit_func *)dlsym(handle, mod_mgr->module_specs[i].instance_exit_cb_name);
+            if (!mod_mgr->module_specs[i].on_instance_exit_cb)
+            {
+                fprintf(stderr, "Could not load exit function %s: %s\n", mod_mgr->module_specs[i].instance_exit_cb_name, dlerror());
+            }
         }
 
-        mod_mgr->module_specs[i].on_exit_cb = (module_on_exit_func *) dlsym(handle, exit_func_name);
-        if (!mod_mgr->module_specs[i].on_exit_cb) {
-            fprintf(stderr, "Could not load exit function %s: %s\n", exit_func_name, dlerror());
+        if (mod_mgr->module_specs[i].thread_init_cb_name)
+        {
+            mod_mgr->module_specs[i].on_thread_init_cb =
+                (module_on_thread_init_func *)dlsym(handle, mod_mgr->module_specs[i].thread_init_cb_name);
+            if (!mod_mgr->module_specs[i].on_thread_init_cb)
+            {
+                fprintf(stderr, "Could not load init function %s: %s\n", mod_mgr->module_specs[i].thread_init_cb_name, dlerror());
+            }
+        }
+
+
+        if (mod_mgr->module_specs[i].thread_exit_cb_name)
+        {
+            mod_mgr->module_specs[i].on_thread_exit_cb =
+                (module_on_thread_exit_func *)dlsym(handle, mod_mgr->module_specs[i].thread_exit_cb_name);
+            if (!mod_mgr->module_specs[i].on_thread_exit_cb)
+            {
+                fprintf(stderr, "Could not load exit function %s: %s\n", mod_mgr->module_specs[i].thread_exit_cb_name, dlerror());
+            }
         }
-		mod_mgr->module_specs[i].path=path;
-		mod_mgr->module_specs[i].init_cb_name=init_func_name;
-		mod_mgr->module_specs[i].exit_cb_name=exit_func_name;
 		mod_mgr->load_module_num+=1;
     }
 
@@ -104,14 +130,16 @@ void stellar_module_manager_free(struct stellar_module_manager *mod_mgr)
     {
         for (int i = 0; i < mod_mgr->load_module_num; i++)
         {
-            if (mod_mgr->module_specs[i].on_exit_cb != NULL &&
-                mod_mgr->module_specs[i].is_init_succ)
+            if (mod_mgr->module_specs[i].on_instance_exit_cb != NULL &&
+                mod_mgr->module_specs[i].init_succ)
             {
-                mod_mgr->module_specs[i].on_exit_cb(mod_mgr, mod_mgr->module_specs[i].mod);
+                mod_mgr->module_specs[i].on_instance_exit_cb(mod_mgr, mod_mgr->module_specs[i].mod);
             }
 			if(mod_mgr->module_specs[i].path)FREE(mod_mgr->module_specs[i].path);
-			if(mod_mgr->module_specs[i].init_cb_name)FREE(mod_mgr->module_specs[i].init_cb_name);
-			if(mod_mgr->module_specs[i].exit_cb_name)FREE(mod_mgr->module_specs[i].exit_cb_name);
+			if(mod_mgr->module_specs[i].instance_init_cb_name)FREE(mod_mgr->module_specs[i].instance_init_cb_name);
+			if(mod_mgr->module_specs[i].instance_exit_cb_name)FREE(mod_mgr->module_specs[i].instance_exit_cb_name);
+			if(mod_mgr->module_specs[i].thread_init_cb_name)FREE(mod_mgr->module_specs[i].thread_init_cb_name);
+			if(mod_mgr->module_specs[i].thread_exit_cb_name)FREE(mod_mgr->module_specs[i].thread_exit_cb_name);
         }
 		FREE(mod_mgr->module_specs);
     }
@@ -156,10 +184,35 @@ struct mq_runtime *stellar_module_manager_get_mq_runtime(struct stellar_module_m
 	return local_mq_rt;
 }
 
-void stellar_module_manager_register_thread(struct stellar_module_manager* mod_mgr __unused, int thread_id, struct mq_runtime *mq_rt)
+void stellar_module_manager_register_thread(struct stellar_module_manager* mod_mgr, int thread_id, struct mq_runtime *mq_rt)
 {
 	local_thread_id=thread_id;
 	local_mq_rt=mq_rt;
+
+	for(int i=0; i<mod_mgr->load_module_num; i++)
+	{
+		if(mod_mgr->module_specs[i].mod == NULL)break;
+		if(mod_mgr->module_specs[i].on_thread_init_cb && mod_mgr->module_specs[i].init_succ)
+		{
+			mod_mgr->module_specs[i].on_thread_init_cb(mod_mgr, thread_id, mod_mgr->module_specs[i].mod);
+		}
+	}
+	return;
+}
+
+void stellar_module_manager_unregister_thread(struct stellar_module_manager *mod_mgr, int thread_id)
+{
+	assert(local_thread_id==thread_id);
+	for(int i=0; i<mod_mgr->load_module_num; i++)
+	{
+		if(mod_mgr->module_specs[i].mod == NULL)break;
+		if(mod_mgr->module_specs[i].on_thread_exit_cb && mod_mgr->module_specs[i].init_succ)
+		{
+			mod_mgr->module_specs[i].on_thread_exit_cb(mod_mgr, thread_id, mod_mgr->module_specs[i].mod);
+		}
+	}
+	local_thread_id=-1;
+	local_mq_rt=NULL;
 	return;
 }
 
@@ -171,7 +224,7 @@ struct stellar_module *stellar_module_manager_get_module(struct stellar_module_m
 		for(int i=0; i<mod_mgr->load_module_num; i++)
 		{
 			if(mod_mgr->module_specs[i].mod == NULL)break;
-			if(strcmp(mod_mgr->module_specs[i].mod->name, module_name)==0 && mod_mgr->module_specs[i].is_init_succ)
+			if(strcmp(mod_mgr->module_specs[i].mod->name, module_name)==0 && mod_mgr->module_specs[i].init_succ)
 			{
 				return mod_mgr->module_specs[i].mod;
 			}
diff --git a/infra/module_manager/module_manager_interna.h b/infra/module_manager/module_manager_interna.h
index b58a374..1fe38f4 100644
--- a/infra/module_manager/module_manager_interna.h
+++ b/infra/module_manager/module_manager_interna.h
@@ -22,12 +22,16 @@ struct stellar_module
 struct module_spec_load
 {
 	struct stellar_module *mod;
-	module_on_init_func *on_init_cb;
-	module_on_exit_func *on_exit_cb;
+	module_on_instance_init_func *on_instance_init_cb;
+	module_on_instance_exit_func *on_instance_exit_cb;
+	module_on_thread_init_func *on_thread_init_cb;
+	module_on_thread_exit_func *on_thread_exit_cb;
     char *path;
-    char *init_cb_name;
-    char *exit_cb_name;
-    bool is_init_succ;
+    char *instance_init_cb_name;
+    char *instance_exit_cb_name;
+    char *thread_init_cb_name;
+    char *thread_exit_cb_name;
+    bool init_succ;
 }__attribute__((aligned(sizeof(void*))));
 
 
diff --git a/infra/module_manager/test/gtest_module_manager_main.cpp b/infra/module_manager/test/gtest_module_manager_main.cpp
index 5a30553..f32c99e 100644
--- a/infra/module_manager/test/gtest_module_manager_main.cpp
+++ b/infra/module_manager/test/gtest_module_manager_main.cpp
@@ -117,6 +117,11 @@ TEST(stellar_module_manager, register_thread) {
 	EXPECT_EQ(stellar_module_manager_get_thread_id(mod_mgr), 1);
 	EXPECT_EQ((long)stellar_module_manager_get_mq_runtime(mod_mgr), 2);
 
+	stellar_module_manager_unregister_thread(mod_mgr, 1);
+
+	EXPECT_EQ(stellar_module_manager_get_thread_id(mod_mgr), -1);
+	EXPECT_EQ((long)stellar_module_manager_get_mq_runtime(mod_mgr), 0);
+
 	stellar_module_manager_free(mod_mgr);	
 
 }
@@ -141,17 +146,31 @@ extern "C" void gtest_module_exit(struct stellar_module_manager *mod_mgr, struct
 
 	EXPECT_EQ(stellar_module_manager_get_module(mod_mgr, "gtest"), mod);
 
-	EXPECT_EQ(stellar_module_manager_get_thread_id(mod_mgr), 1);
-	EXPECT_EQ((long)stellar_module_manager_get_mq_runtime(mod_mgr), 2);
+	EXPECT_EQ(stellar_module_manager_get_thread_id(mod_mgr), -1);
+	EXPECT_EQ((long)stellar_module_manager_get_mq_runtime(mod_mgr), 0);
 
 	stellar_module_free(mod);
 }
 
+extern "C" void gtest_thread_init(struct stellar_module_manager *mod_mgr, int thread_id, struct stellar_module *mod)
+{
+	EXPECT_STREQ(stellar_module_get_name(mod), "gtest");
+	EXPECT_EQ((long)stellar_module_get_ctx(mod), 1);
+}
+
+extern "C" void gtest_thread_exit(struct stellar_module_manager *mod_mgr, int thread_id, struct stellar_module *mod)
+{
+	EXPECT_STREQ(stellar_module_get_name(mod), "gtest");
+	EXPECT_EQ((long)stellar_module_get_ctx(mod), 1);
+}
+
 const char *gtest_module_spec_toml = 
     "[[module]]\n"
     "path = \"\"\n"
     "init = \"gtest_module_init\"\n"
-    "exit = \"gtest_module_exit\"\n";
+    "exit = \"gtest_module_exit\"\n"
+	"thread_init = \"gtest_thread_init\"\n"
+    "thread_exit = \"gtest_thread_exit\"\n";
 
 TEST(module_manager, basic_module) {
 
@@ -178,6 +197,10 @@ TEST(module_manager, basic_module) {
 	EXPECT_EQ((long)stellar_module_manager_get_thread_id(mod_mgr), 1);
 	EXPECT_EQ((long)stellar_module_manager_get_mq_runtime(mod_mgr), 2);
 
+	stellar_module_manager_unregister_thread(mod_mgr, 1);
+	EXPECT_EQ((long)stellar_module_manager_get_thread_id(mod_mgr), -1);
+	EXPECT_EQ((long)stellar_module_manager_get_mq_runtime(mod_mgr), 0);
+
 	stellar_module_manager_free(mod_mgr);	
 	unlink(toml_template);
 }
@@ -205,7 +228,7 @@ struct test_module_polling_env
  }
 
 
-TEST(polling_manager, basic_polling_module) {
+TEST(module_manager, basic_polling_module) {
 
 	struct mq_schema *mq_schema=mq_schema_new();
 
@@ -232,6 +255,10 @@ TEST(polling_manager, basic_polling_module) {
 		stellar_polling_dispatch(mod_mgr);
 	}
 
+	stellar_module_manager_unregister_thread(mod_mgr, 1);
+
+	mq_runtime_free(mq_rt);
+	mq_schema_free(mq_schema);
 	stellar_module_manager_free(mod_mgr);	
 	
 	EXPECT_EQ(env.polling_count, env.N_round+env.polling_active_count);
diff --git a/infra/stellar_core.c b/infra/stellar_core.c
index bfed599..8c6e704 100644
--- a/infra/stellar_core.c
+++ b/infra/stellar_core.c
@@ -115,7 +115,7 @@ static void *worker_thread(void *arg)
             packet_io_yield(pkt_io, thread_id);
         }
     }
-
+    stellar_module_manager_unregister_thread(mod_mgr, thread_id);
     mq_runtime_free(mq_rt);
 
     ATOMIC_SET(&thread->is_runing, 0);