From 7d7224a6a97294f4d9ce39a08658e58950d8f176 Mon Sep 17 00:00:00 2001
From: Sean McBride <spmcbride1201@gmail.com>
Date: Mon, 9 May 2022 21:42:04 -0400
Subject: [PATCH] refactor: module pools

---
 runtime/include/module.h | 18 ++++++++++++------
 runtime/src/main.c       |  2 --
 runtime/src/module.c     | 22 +++++++++++++++++-----
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/runtime/include/module.h b/runtime/include/module.h
index 4a19cd5..94f3f5a 100644
--- a/runtime/include/module.h
+++ b/runtime/include/module.h
@@ -23,10 +23,7 @@ extern thread_local int worker_thread_idx;
 INIT_POOL(wasm_memory, wasm_memory_free)
 INIT_POOL(wasm_stack, wasm_stack_free)
 
-/* TODO: Dynamically size based on number of threads */
-#define MAX_WORKER_THREADS 64
-
-struct module_pools {
+struct module_pool {
 	struct wasm_memory_pool memory;
 	struct wasm_stack_pool  stack;
 } __attribute__((aligned(CACHE_PAD)));
@@ -41,7 +38,7 @@ struct module {
 	_Atomic uint32_t               reference_count; /* ref count how many instances exist here. */
 	struct sledge_abi__wasm_table *indirect_table;
 
-	struct module_pools pools[MAX_WORKER_THREADS];
+	struct module_pool *pools;
 } __attribute__((aligned(CACHE_PAD)));
 
 /********************************
@@ -118,12 +115,21 @@ module_alloc_table(struct module *module)
 static inline void
 module_initialize_pools(struct module *module)
 {
-	for (int i = 0; i < MAX_WORKER_THREADS; i++) {
+	for (int i = 0; i < runtime_worker_threads_count; i++) {
 		wasm_memory_pool_init(&module->pools[i].memory, false);
 		wasm_stack_pool_init(&module->pools[i].stack, false);
 	}
 }
 
+static inline void
+module_deinitialize_pools(struct module *module)
+{
+	for (int i = 0; i < runtime_worker_threads_count; i++) {
+		wasm_memory_pool_deinit(&module->pools[i].memory);
+		wasm_stack_pool_deinit(&module->pools[i].stack);
+	}
+}
+
 /**
  * Invoke a module's initialize_memory
  * @param module - the module whose memory we are initializing
diff --git a/runtime/src/main.c b/runtime/src/main.c
index 1c12a59..902ff46 100644
--- a/runtime/src/main.c
+++ b/runtime/src/main.c
@@ -35,8 +35,6 @@ uint32_t runtime_total_online_processors = 0;
 uint32_t runtime_worker_threads_count    = 0;
 
 enum RUNTIME_SIGALRM_HANDLER runtime_sigalrm_handler = RUNTIME_SIGALRM_HANDLER_BROADCAST;
-int                          runtime_worker_core_count;
-
 
 bool     runtime_preemption_enabled = true;
 uint32_t runtime_quantum_us         = 5000; /* 5ms */
diff --git a/runtime/src/module.c b/runtime/src/module.c
index bb69bf6..69cebfd 100644
--- a/runtime/src/module.c
+++ b/runtime/src/module.c
@@ -39,6 +39,8 @@ module_init(struct module *module, IN char *path)
 	rc = sledge_abi_symbols_init(&module->abi, path);
 	if (rc != 0) goto err;
 
+	module->pools = calloc(runtime_worker_threads_count, sizeof(struct module_pool));
+
 	module->path = path;
 
 	module->stack_size = ((uint32_t)(round_up_to_page(stack_size == 0 ? WASM_STACK_SIZE : stack_size)));
@@ -52,6 +54,19 @@ err:
 	goto done;
 }
 
+static inline void
+module_deinit(struct module *module)
+{
+	assert(module == NULL);
+	assert(module->reference_count == 0);
+
+	free(module->path);
+	sledge_abi_symbols_deinit(&module->abi);
+	/* TODO: Free indirect_table */
+	module_deinitialize_pools(module);
+	free(module->pools);
+}
+
 /***************************************
  * Public Methods
  ***************************************/
@@ -69,14 +84,11 @@ void
 module_free(struct module *module)
 {
 	if (module == NULL) return;
+	assert(module->reference_count == 0);
 
 	panic("Unimplemented!\n");
 
-	/* TODO: Should allocating routes increment reference */
-	/* Do not free if we still have oustanding references */
-	if (module->reference_count) return;
-
-	sledge_abi_symbols_deinit(&module->abi);
+	module_deinit(module);
 	free(module);
 }