Enable aux stack allocations on application heap (#1799)

This is necessary for WASI threads as the aux stack should be managed by the application.
See https://github.com/bytecodealliance/wasm-micro-runtime/issues/1790 for details.

core/config.h

@@ -165,6 +165,13 @@
 #define WASM_ENABLE_LIB_WASI_THREADS 0
 #endif
 
+#ifndef WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION
+#define WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION WASM_ENABLE_LIB_WASI_THREADS
+#elif WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION == 0 \
+    && WASM_ENABLE_LIB_WASI_THREADS == 1
+#error "Heap aux stack allocation must be enabled for WASI threads"
+#endif
+
 #ifndef WASM_ENABLE_BASE_LIB
 #define WASM_ENABLE_BASE_LIB 0
 #endif

core/iwasm/interpreter/wasm_runtime.c

@@ -2443,14 +2443,16 @@ wasm_set_aux_stack(WASMExecEnv *exec_env, uint32 start_offset, uint32 size)
     WASMModuleInstance *module_inst =
         (WASMModuleInstance *)exec_env->module_inst;
     uint32 stack_top_idx = module_inst->module->aux_stack_top_global_index;
+
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION == 0
+    /* Check the aux stack space */
     uint32 data_end = module_inst->module->aux_data_end;
     uint32 stack_bottom = module_inst->module->aux_stack_bottom;
     bool is_stack_before_data = stack_bottom < data_end ? true : false;
-
-    /* Check the aux stack space, currently we don't allocate space in heap */
     if ((is_stack_before_data && (size > start_offset))
         || ((!is_stack_before_data) && (start_offset - data_end < size)))
         return false;
+#endif
 
     if (stack_top_idx != (uint32)-1) {
         /* The aux stack top is a wasm global,

core/iwasm/libraries/lib-pthread/lib_pthread_wrapper.c

@@ -504,7 +504,6 @@ pthread_start_routine(void *arg)
     info_node->exec_env = exec_env;
     info_node->u.thread = exec_env->handle;
     if (!append_thread_info_node(info_node)) {
-        wasm_runtime_deinstantiate_internal(module_inst, true);
         delete_thread_info_node(info_node);
         os_cond_signal(&parent_exec_env->wait_cond);
         os_mutex_unlock(&parent_exec_env->wait_lock);
@@ -527,9 +526,6 @@ pthread_start_routine(void *arg)
     /* destroy pthread key values */
     call_key_destructor(exec_env);
 
-    /* routine exit, destroy instance */
-    wasm_runtime_deinstantiate_internal(module_inst, true);
-
     wasm_runtime_free(routine_args);
 
     /* if the thread is joinable, store the result in its info node,

core/iwasm/libraries/lib-wasi-threads/lib_wasi_threads.cmake

@@ -3,7 +3,7 @@
 
 set (LIB_WASI_THREADS_DIR ${CMAKE_CURRENT_LIST_DIR})
 
-add_definitions (-DWASM_ENABLE_LIB_WASI_THREADS=1)
+add_definitions (-DWASM_ENABLE_LIB_WASI_THREADS=1 -DWASM_ENABLE_HEAP_AUX_STACK_ALLOCATION=1)
 
 include_directories(${LIB_WASI_THREADS_DIR})
 

core/iwasm/libraries/lib-wasi-threads/lib_wasi_threads_wrapper.c

@@ -4,7 +4,6 @@
  */
 
 #include "bh_log.h"
-#include "wasmtime_ssp.h"
 #include "thread_manager.h"
 
 #if WASM_ENABLE_INTERP != 0
@@ -58,9 +57,6 @@ thread_start(void *arg)
             wasm_cluster_spread_exception(exec_env);
     }
 
-    /* routine exit, destroy instance */
-    wasm_runtime_deinstantiate_internal(module_inst, true);
-
     wasm_runtime_free(thread_arg);
     exec_env->thread_arg = NULL;
 

core/iwasm/libraries/thread-mgr/thread_manager.c

@@ -77,8 +77,18 @@ traverse_list(bh_list *l, list_visitor visitor, void *user_data)
 }
 
 static bool
-allocate_aux_stack(WASMCluster *cluster, uint32 *start, uint32 *size)
+allocate_aux_stack(WASMExecEnv *exec_env, uint32 *start, uint32 *size)
 {
+    WASMCluster *cluster = wasm_exec_env_get_cluster(exec_env);
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION != 0
+    WASMModuleInstanceCommon *module_inst =
+        wasm_exec_env_get_module_inst(exec_env);
+
+    *start = wasm_runtime_module_malloc(module_inst, cluster->stack_size, NULL);
+    *size = cluster->stack_size;
+
+    return *start != 0;
+#else
     uint32 i;
 
     /* If the module doesn't have aux stack info,
@@ -100,11 +110,21 @@ allocate_aux_stack(WASMCluster *cluster, uint32 *start, uint32 *size)
     }
     os_mutex_unlock(&cluster->lock);
     return false;
+#endif
 }
 
 static bool
-free_aux_stack(WASMCluster *cluster, uint32 start)
+free_aux_stack(WASMExecEnv *exec_env, uint32 start)
 {
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION != 0
+    WASMModuleInstanceCommon *module_inst =
+        wasm_exec_env_get_module_inst(exec_env);
+
+    wasm_runtime_module_free(module_inst, start);
+
+    return true;
+#else
+    WASMCluster *cluster = wasm_exec_env_get_cluster(exec_env);
     uint32 i;
 
     for (i = 0; i < cluster_max_thread_num; i++) {
@@ -116,14 +136,14 @@ free_aux_stack(WASMCluster *cluster, uint32 start)
         }
     }
     return false;
+#endif
 }
 
 WASMCluster *
 wasm_cluster_create(WASMExecEnv *exec_env)
 {
     WASMCluster *cluster;
-    uint64 total_size;
+    uint32 aux_stack_start, aux_stack_size;
-    uint32 aux_stack_start, aux_stack_size, i;
 
     bh_assert(exec_env->cluster == NULL);
     if (!(cluster = wasm_runtime_malloc(sizeof(WASMCluster)))) {
@@ -159,12 +179,16 @@ wasm_cluster_create(WASMExecEnv *exec_env)
         return cluster;
     }
 
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION != 0
+    cluster->stack_size = aux_stack_size;
+#else
     cluster->stack_size = aux_stack_size / (cluster_max_thread_num + 1);
     if (cluster->stack_size < WASM_THREAD_AUX_STACK_SIZE_MIN) {
         goto fail;
     }
     /* Make stack size 16-byte aligned */
     cluster->stack_size = cluster->stack_size & (~15);
+#endif
 
     /* Set initial aux stack top to the instance and
         aux stack boundary to the main exec_env */
@@ -172,8 +196,10 @@ wasm_cluster_create(WASMExecEnv *exec_env)
                                      cluster->stack_size))
         goto fail;
 
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION == 0
     if (cluster_max_thread_num != 0) {
-        total_size = cluster_max_thread_num * sizeof(uint32);
+        uint64 total_size = cluster_max_thread_num * sizeof(uint32);
+        uint32 i;
         if (total_size >= UINT32_MAX
             || !(cluster->stack_tops =
                      wasm_runtime_malloc((uint32)total_size))) {
@@ -195,6 +221,7 @@ wasm_cluster_create(WASMExecEnv *exec_env)
             cluster->stack_tops[i] = aux_stack_start - cluster->stack_size * i;
         }
     }
+#endif
 
     os_mutex_lock(&cluster_list_lock);
     if (bh_list_insert(cluster_list, cluster) != 0) {
@@ -234,10 +261,12 @@ wasm_cluster_destroy(WASMCluster *cluster)
 
     os_mutex_destroy(&cluster->lock);
 
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION == 0
     if (cluster->stack_tops)
         wasm_runtime_free(cluster->stack_tops);
     if (cluster->stack_segment_occupied)
         wasm_runtime_free(cluster->stack_segment_occupied);
+#endif
 
 #if WASM_ENABLE_DEBUG_INTERP != 0
     wasm_debug_instance_destroy(cluster);
@@ -273,7 +302,13 @@ wasm_cluster_add_exec_env(WASMCluster *cluster, WASMExecEnv *exec_env)
     exec_env->cluster = cluster;
 
     os_mutex_lock(&cluster->lock);
-    if (bh_list_insert(&cluster->exec_env_list, exec_env) != 0)
+    if (cluster->exec_env_list.len == cluster_max_thread_num + 1) {
+        LOG_ERROR("thread manager error: "
+                  "maximum number of threads exceeded");
+        ret = false;
+    }
+
+    if (ret && bh_list_insert(&cluster->exec_env_list, exec_env) != 0)
         ret = false;
     os_mutex_unlock(&cluster->lock);
     return ret;
@@ -407,7 +442,7 @@ wasm_cluster_spawn_exec_env(WASMExecEnv *exec_env)
     if (!new_exec_env)
         goto fail1;
 
-    if (!allocate_aux_stack(cluster, &aux_stack_start, &aux_stack_size)) {
+    if (!allocate_aux_stack(exec_env, &aux_stack_start, &aux_stack_size)) {
         LOG_ERROR("thread manager error: "
                   "failed to allocate aux stack space for new thread");
         goto fail2;
@@ -426,7 +461,7 @@ wasm_cluster_spawn_exec_env(WASMExecEnv *exec_env)
 
 fail3:
     /* free the allocated aux stack space */
-    free_aux_stack(cluster, aux_stack_start);
+    free_aux_stack(exec_env, aux_stack_start);
 fail2:
     wasm_exec_env_destroy(new_exec_env);
 fail1:
@@ -443,7 +478,7 @@ wasm_cluster_destroy_spawned_exec_env(WASMExecEnv *exec_env)
     bh_assert(cluster != NULL);
 
     /* Free aux stack space */
-    free_aux_stack(cluster, exec_env->aux_stack_bottom.bottom);
+    free_aux_stack(exec_env, exec_env->aux_stack_bottom.bottom);
     wasm_cluster_del_exec_env(cluster, exec_env);
     wasm_exec_env_destroy_internal(exec_env);
 
@@ -457,7 +492,11 @@ thread_manager_start_routine(void *arg)
     void *ret;
     WASMExecEnv *exec_env = (WASMExecEnv *)arg;
     WASMCluster *cluster = wasm_exec_env_get_cluster(exec_env);
+    WASMModuleInstanceCommon *module_inst =
+        wasm_exec_env_get_module_inst(exec_env);
+
     bh_assert(cluster != NULL);
+    bh_assert(module_inst != NULL);
 
     exec_env->handle = os_self_thread();
     ret = exec_env->thread_start_routine(exec_env);
@@ -469,7 +508,11 @@ thread_manager_start_routine(void *arg)
 
     /* Routine exit */
     /* Free aux stack space */
-    free_aux_stack(cluster, exec_env->aux_stack_bottom.bottom);
+    free_aux_stack(exec_env, exec_env->aux_stack_bottom.bottom);
+
+    /* routine exit, destroy instance */
+    wasm_runtime_deinstantiate_internal(module_inst, true);
+
     /* Detach the native thread here to ensure the resources are freed */
     wasm_cluster_detach_thread(exec_env);
 #if WASM_ENABLE_DEBUG_INTERP != 0
@@ -501,7 +544,7 @@ wasm_cluster_create_thread(WASMExecEnv *exec_env,
     if (!new_exec_env)
         return -1;
 
-    if (!allocate_aux_stack(cluster, &aux_stack_start, &aux_stack_size)) {
+    if (!allocate_aux_stack(exec_env, &aux_stack_start, &aux_stack_size)) {
         LOG_ERROR("thread manager error: "
                   "failed to allocate aux stack space for new thread");
         goto fail1;
@@ -532,7 +575,7 @@ fail3:
     wasm_cluster_del_exec_env(cluster, new_exec_env);
 fail2:
     /* free the allocated aux stack space */
-    free_aux_stack(cluster, aux_stack_start);
+    free_aux_stack(exec_env, aux_stack_start);
 fail1:
     wasm_exec_env_destroy(new_exec_env);
     return -1;
@@ -762,7 +805,7 @@ wasm_cluster_exit_thread(WASMExecEnv *exec_env, void *retval)
 #endif
     /* App exit the thread, free the resources before exit native thread */
     /* Free aux stack space */
-    free_aux_stack(cluster, exec_env->aux_stack_bottom.bottom);
+    free_aux_stack(exec_env, exec_env->aux_stack_bottom.bottom);
     /* Detach the native thread here to ensure the resources are freed */
     wasm_cluster_detach_thread(exec_env);
     /* Remove and destroy exec_env */

core/iwasm/libraries/thread-mgr/thread_manager.h

@@ -26,14 +26,16 @@ struct WASMCluster {
     korp_mutex lock;
     bh_list exec_env_list;
 
+#if WASM_ENABLE_HEAP_AUX_STACK_ALLOCATION == 0
     /* The aux stack of a module with shared memory will be
         divided into several segments. This array store the
         stack top of different segments */
     uint32 *stack_tops;
-    /* Size of every stack segment */
-    uint32 stack_size;
     /* Record which segments are occupied */
     bool *stack_segment_occupied;
+#endif
+    /* Size of every stack segment */
+    uint32 stack_size;
 #if WASM_ENABLE_DEBUG_INTERP != 0
     WASMDebugInstance *debug_inst;
 #endif