From 534a8cf9f4b812950155abf17d98430e1d5599e9 Mon Sep 17 00:00:00 2001
From: YAMAMOTO Takashi <yamamoto@midokura.com>
Date: Fri, 8 Sep 2023 12:49:04 +0900
Subject: [PATCH] Handle a return from wasi _start function correctly (#2529)

This fixes a few test cases in wasi-threads testsuite like wasi_threads_return_main_block.
And also move the special handling for "wasi proc exit" to a more appropriate place.
---
 core/iwasm/common/wasm_application.c    | 29 ++++++++++++++++-
 core/iwasm/common/wasm_runtime_common.c | 42 ++-----------------------
 2 files changed, 31 insertions(+), 40 deletions(-)

diff --git a/core/iwasm/common/wasm_application.c b/core/iwasm/common/wasm_application.c
index 2ed217e7a..795d65c8c 100644
--- a/core/iwasm/common/wasm_application.c
+++ b/core/iwasm/common/wasm_application.c
@@ -107,7 +107,34 @@ execute_main(WASMModuleInstanceCommon *module_inst, int32 argc, char *argv[])
        the actual main function. Directly calling main function
        may cause exception thrown. */
     if ((func = wasm_runtime_lookup_wasi_start_function(module_inst))) {
-        return wasm_runtime_call_wasm(exec_env, func, 0, NULL);
+        const char *wasi_proc_exit_exception = "wasi proc exit";
+
+        ret = wasm_runtime_call_wasm(exec_env, func, 0, NULL);
+#if WASM_ENABLE_THREAD_MGR != 0
+        if (ret) {
+            /* On a successful return from the `_start` function,
+               we terminate other threads by mimicing wasi:proc_exit(0).
+
+               Note:
+               - A return from the `main` function is an equivalent of
+                 exit(). (C standard)
+               - When exit code is 0, wasi-libc's `_start` function just
+                 returns w/o calling `proc_exit`.
+               - A process termination should terminate threads in
+                 the process. */
+
+            wasm_runtime_set_exception(module_inst, wasi_proc_exit_exception);
+            /* exit_code is zero-initialized */
+            ret = false;
+        }
+#endif
+        /* report wasm proc exit as a success */
+        WASMModuleInstance *inst = (WASMModuleInstance *)module_inst;
+        if (!ret && strstr(inst->cur_exception, wasi_proc_exit_exception)) {
+            inst->cur_exception[0] = 0;
+            ret = true;
+        }
+        return ret;
     }
 #endif /* end of WASM_ENABLE_LIBC_WASI */
 
diff --git a/core/iwasm/common/wasm_runtime_common.c b/core/iwasm/common/wasm_runtime_common.c
index cb7dbee9d..7cac9e643 100644
--- a/core/iwasm/common/wasm_runtime_common.c
+++ b/core/iwasm/common/wasm_runtime_common.c
@@ -1938,33 +1938,6 @@ wasm_runtime_finalize_call_function(WASMExecEnv *exec_env,
 }
 #endif
 
-static bool
-clear_wasi_proc_exit_exception(WASMModuleInstanceCommon *module_inst_comm)
-{
-#if WASM_ENABLE_LIBC_WASI != 0
-    bool has_exception;
-    char exception[EXCEPTION_BUF_LEN];
-    WASMModuleInstance *module_inst = (WASMModuleInstance *)module_inst_comm;
-
-    bh_assert(module_inst_comm->module_type == Wasm_Module_Bytecode
-              || module_inst_comm->module_type == Wasm_Module_AoT);
-
-    has_exception = wasm_copy_exception(module_inst, exception);
-    if (has_exception && !strcmp(exception, "Exception: wasi proc exit")) {
-        /* The "wasi proc exit" exception is thrown by native lib to
-           let wasm app exit, which is a normal behavior, we clear
-           the exception here. And just clear the exception of current
-           thread, don't call `wasm_set_exception(module_inst, NULL)`
-           which will clear the exception of all threads. */
-        module_inst->cur_exception[0] = '\0';
-        return true;
-    }
-    return false;
-#else
-    return false;
-#endif
-}
-
 bool
 wasm_runtime_call_wasm(WASMExecEnv *exec_env,
                        WASMFunctionInstanceCommon *function, uint32 argc,
@@ -2005,15 +1978,10 @@ wasm_runtime_call_wasm(WASMExecEnv *exec_env,
                                 param_argc, new_argv);
 #endif
     if (!ret) {
-        if (clear_wasi_proc_exit_exception(exec_env->module_inst)) {
-            ret = true;
-        }
-        else {
-            if (new_argv != argv) {
-                wasm_runtime_free(new_argv);
-            }
-            return false;
+        if (new_argv != argv) {
+            wasm_runtime_free(new_argv);
         }
+        return false;
     }
 
 #if WASM_ENABLE_REF_TYPES != 0
@@ -4573,10 +4541,6 @@ wasm_runtime_call_indirect(WASMExecEnv *exec_env, uint32 element_index,
         ret = aot_call_indirect(exec_env, 0, element_index, argc, argv);
 #endif
 
-    if (!ret && clear_wasi_proc_exit_exception(exec_env->module_inst)) {
-        ret = true;
-    }
-
     return ret;
 }