Enhance wasm with checkpoint and restore support (#2333)

ATTRIBUTIONS.md

@@ -18,7 +18,7 @@ WAMR project reused some components from other open source project:
 - **NuttX ELF headers**: used in core/iwasm/aot/debug/elf_parser.c
 - **Dhrystone**: for the test benchmakr dhrystone
 
-The WAMR fast interpreter is a clean room development. We would acknowledge the inspirations by [WASM3](https://github.com/wasm3/wasm3) open source project for the approach of pre-calculated oprand stack location.
+The WAMR fast interpreter is a clean room development. We would acknowledge the inspirations by [WASM3](https://github.com/wasm3/wasm3) open source project for the approach of pre-calculated operand stack location.
 
 |  third party components | version number | latest release | vendor pages | CVE details |
 | --- | --- | --- | --- | --- |
@@ -37,6 +37,7 @@ The WAMR fast interpreter is a clean room development. We would acknowledge the
 | zydis | unspecified | e14a07895136182a5b53e181eec3b1c6e0b434de | https://github.com/zyantific/zydis | |
 | NuttX ELF headers | 72313301e23f9c2de969fb64b9a0f67bb4c284df | 10.3.0 | https://github.com/apache/nuttx | |
 | Dhrystone | 2.1 | 2.1 | https://fossies.org/linux/privat/old/ | |
+| yalantinglibs | v0.3.2 | v0.3.2 | https://github.com/alibaba/yalantinglibs | |
 
 ## Licenses
 
@@ -103,3 +104,7 @@ The WAMR fast interpreter is a clean room development. We would acknowledge the
 ### Dhrystone
 
 [LICENSE](./tests/benchmarks/dhrystone/LICENSE)
+
+### yalantinglibs
+
+[LICENSE](./core/iwasm/libraries/ckpt-restore/LICENSE_YALANTINGLIBS)

build-scripts/config_common.cmake

@@ -315,6 +315,10 @@ if (WAMR_BUILD_AOT_STACK_FRAME EQUAL 1)
   add_definitions (-DWASM_ENABLE_AOT_STACK_FRAME=1)
   message ("     AOT stack frame enabled")
 endif ()
+if (WAMR_BUILD_CHECKPOINT_RESTORE EQUAL 1)
+  add_definitions (-DWASM_ENABLE_CHECKPOINT_RESTORE=1)
+  message ("     Checkpoint Restore enabled")
+endif ()
 if (WAMR_BUILD_MEMORY_PROFILING EQUAL 1)
   add_definitions (-DWASM_ENABLE_MEMORY_PROFILING=1)
   message ("     Memory profiling enabled")

build-scripts/runtime_lib.cmake

@@ -86,6 +86,10 @@ if (WAMR_BUILD_LIBC_BUILTIN EQUAL 1)
     include (${IWASM_DIR}/libraries/libc-builtin/libc_builtin.cmake)
 endif ()
 
+if (WAMR_BUILD_CHECKPOINT_RESTORE EQUAL 1)
+    include (${IWASM_DIR}/libraries/ckpt-restore/ckpt_restore.cmake)
+endif ()
+
 if (WAMR_BUILD_LIBC_UVWASI EQUAL 1)
     include (${IWASM_DIR}/libraries/libc-uvwasi/libc_uvwasi.cmake)
     set (WAMR_BUILD_MODULE_INST_CONTEXT 1)
@@ -193,6 +197,7 @@ set (source_all
     ${LIBC_EMCC_SOURCE}
     ${LIB_RATS_SOURCE}
     ${DEBUG_ENGINE_SOURCE}
+    ${CKPT_RESTORE_SOURCE}
 )
 
 set (WAMR_RUNTIME_LIB_SOURCE ${source_all})

core/config.h

@@ -339,6 +339,11 @@
 #define WASM_ENABLE_AOT_STACK_FRAME 0
 #endif
 
+/* Checkpoint Restore */
+#ifndef WASM_ENABLE_CHECKPOINT_RESTORE
+#define WASM_ENABLE_CHECKPOINT_RESTORE 0
+#endif
+
 /* Heap verification */
 #ifndef BH_ENABLE_GC_VERIFY
 #define BH_ENABLE_GC_VERIFY 0

core/iwasm/aot/aot_reloc.h

@@ -57,6 +57,7 @@ typedef struct {
 #define REG_AOT_TRACE_SYM()               \
     REG_SYM(aot_alloc_frame),             \
     REG_SYM(aot_free_frame),              \
+    REG_SYM(aot_raise),                   \
     REG_SYM(aot_frame_update_profile_info),
 #else
 #define REG_AOT_TRACE_SYM()
@@ -246,4 +247,4 @@ apply_relocation(AOTModule *module,
 }
 #endif
 
-#endif /* end of _AOT_RELOC_H_ */
+#endif /* end of _AOT_RELOC_H_ */

core/iwasm/aot/aot_runtime.c

@@ -15,6 +15,9 @@
 #if WASM_ENABLE_THREAD_MGR != 0
 #include "../libraries/thread-mgr/thread_manager.h"
 #endif
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+#include "../libraries/ckpt-restore/ckpt_restore.h"
+#endif
 
 /*
  * Note: These offsets need to match the values hardcoded in
@@ -2857,6 +2860,15 @@ aot_call_indirect(WASMExecEnv *exec_env, uint32 tbl_idx, uint32 table_elem_idx,
     }
 
     tbl_elem_val = ((table_elem_type_t *)tbl_inst->elems)[table_elem_idx];
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    if (exec_env->is_restore && exec_env->restore_call_chain) {
+        struct AOTFrame *rcc = *(exec_env->restore_call_chain);
+        while (rcc->prev_frame) {
+            rcc = rcc->prev_frame;
+        }
+        func_idx = rcc->func_index;
+    }
+#endif
     if (tbl_elem_val == NULL_REF) {
         aot_set_exception_with_id(module_inst, EXCE_UNINITIALIZED_ELEMENT);
         goto fail;
@@ -3529,6 +3541,12 @@ get_func_name_from_index(const AOTModuleInstance *module_inst,
 #endif /* end of WASM_ENABLE_DUMP_CALL_STACK != 0 || \
           WASM_ENABLE_PERF_PROFILING != 0 */
 
+void
+aot_raise(WASMExecEnv *exec_env, int sig)
+{
+    raise(sig);
+}
+
 #if WASM_ENABLE_GC == 0
 bool
 aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
@@ -3608,6 +3626,9 @@ aot_free_frame(WASMExecEnv *exec_env)
 bool
 aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    LOG_DEBUG("aot_alloc_frame %u thread %d\n", func_index, exec_env->handle);
+#endif
     AOTModuleInstance *module_inst = (AOTModuleInstance *)exec_env->module_inst;
     AOTModule *module = (AOTModule *)module_inst->module;
 #if WASM_ENABLE_PERF_PROFILING != 0
@@ -3617,6 +3638,27 @@ aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
     AOTFrame *frame;
     uint32 max_local_cell_num, max_stack_cell_num, all_cell_num;
     uint32 aot_func_idx, frame_size;
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    if (exec_env->restore_call_chain) {
+        frame = exec_env->restore_call_chain[exec_env->call_chain_size - 1];
+        LOG_DEBUG("frame restored, func idx %zu\n", frame->func_index);
+        exec_env->call_chain_size--;
+        frame->prev_frame = (AOTFrame *)exec_env->cur_frame;
+        exec_env->cur_frame = (struct WASMInterpFrame *)frame;
+        if (exec_env->call_chain_size == 0) {
+            // TODO: fix memory leak
+            exec_env->restore_call_chain = NULL;
+        }
+        LOG_DEBUG("restore call chain %zu==%u, %p, %p, %d\n",
+                  ((AOTFrame *)exec_env->cur_frame)->func_index, func_index,
+                  exec_env, exec_env->restore_call_chain, exec_env->handle);
+        if (((AOTFrame *)exec_env->cur_frame)->func_index != func_index) {
+            LOG_DEBUG("NOT MATCH!!!\n");
+            exit(1);
+        }
+        return true;
+    }
+#endif
 
     if (func_index >= module->import_func_count) {
         aot_func_idx = func_index - module->import_func_count;
@@ -3648,6 +3690,11 @@ aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
     frame->time_started = (uintptr_t)os_time_thread_cputime_us();
     frame->func_perf_prof_info = func_perf_prof;
 #endif
+    frame->ip_offset = 0;
+    frame->sp = frame->lp + max_local_cell_num;
+#if WASM_ENABLE_GC != 0
+    frame->frame_ref = frame->sp + max_stack_cell_num;
+#endif
 
 #if WASM_ENABLE_GC != 0
     frame->sp = frame->lp + max_local_cell_num;
@@ -3664,6 +3711,10 @@ aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
 static inline void
 aot_free_frame_internal(WASMExecEnv *exec_env)
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    int func_index = ((AOTFrame *)exec_env->cur_frame)->func_index;
+    LOG_DEBUG("aot_free_frame %zu %d\n", func_index, exec_env->handle);
+#endif
     AOTFrame *cur_frame = (AOTFrame *)exec_env->cur_frame;
     AOTFrame *prev_frame = cur_frame->prev_frame;
 

core/iwasm/aot/aot_runtime.h

@@ -678,6 +678,9 @@ aot_table_grow(AOTModuleInstance *module_inst, uint32 tbl_idx,
 bool
 aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index);
 
+void
+aot_raise(WASMExecEnv *exec_env, int exception);
+
 void
 aot_free_frame(WASMExecEnv *exec_env);
 

core/iwasm/common/wasm_exec_env.c

@@ -72,6 +72,9 @@ wasm_exec_env_create_internal(struct WASMModuleInstanceCommon *module_inst,
     exec_env->wasm_stack.top_boundary =
         exec_env->wasm_stack.bottom + stack_size;
     exec_env->wasm_stack.top = exec_env->wasm_stack.bottom;
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    exec_env->is_checkpoint = false;
+#endif
 
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT) {
@@ -85,6 +88,12 @@ wasm_exec_env_create_internal(struct WASMModuleInstanceCommon *module_inst,
     wasm_runtime_dump_exec_env_mem_consumption(exec_env);
 #endif
 
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    exec_env->is_checkpoint = false;
+    exec_env->is_restore = false;
+    exec_env->call_chain_size = 0;
+    exec_env->restore_call_chain = NULL;
+#endif
     return exec_env;
 
 #ifdef OS_ENABLE_HW_BOUND_CHECK

core/iwasm/common/wasm_exec_env.h

@@ -155,7 +155,14 @@ typedef struct WASMExecEnv {
 
     /* The WASM stack size */
     uint32 wasm_stack_size;
-
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    /* Whether is checkpoint */
+    bool is_checkpoint;
+    /* Whether is restore */
+    bool is_restore;
+    size_t call_chain_size;
+    struct AOTFrame **restore_call_chain;
+#endif
     /* The WASM stack of current thread */
     union {
         uint64 __make_it_8_byte_aligned_;

core/iwasm/common/wasm_memory.c

@@ -585,7 +585,6 @@ wasm_runtime_get_native_addr_range(WASMModuleInstanceCommon *module_inst_comm,
     SHARED_MEMORY_UNLOCK(memory_inst);
     return false;
 }
-
 bool
 wasm_check_app_addr_and_convert(WASMModuleInstance *module_inst, bool is_str,
                                 uint64 app_buf_addr, uint64 app_buf_size,

core/iwasm/common/wasm_runtime_common.c

@@ -37,6 +37,9 @@
 #if WASM_ENABLE_JIT != 0 || WASM_ENABLE_WAMR_COMPILER != 0
 #include "../compilation/aot_llvm.h"
 #endif
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+#include "../libraries/ckpt-restore/ckpt_restore.h"
+#endif
 #include "../common/wasm_c_api_internal.h"
 #include "../../version.h"
 
@@ -62,6 +65,10 @@
 #undef CHECK
 #undef CHECK1
 
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+#undef wasm_runtime_invoke_native
+#endif
+
 #if WASM_ENABLE_MULTI_MODULE != 0
 /**
  * A safety insurance to prevent
@@ -3570,7 +3577,8 @@ wasm_runtime_init_wasi(WASMModuleInstanceCommon *module_inst,
     uint32 i;
     bool ret = false;
 
-    ctx = runtime_malloc(sizeof(*ctx), module_inst, error_buf, error_buf_size);
+    ctx = runtime_malloc(sizeof(WASIContext), module_inst, error_buf,
+                         error_buf_size);
     if (!ctx)
         return false;
     uvwasi = &ctx->uvwasi;
@@ -5701,6 +5709,9 @@ bool
 wasm_runtime_call_indirect(WASMExecEnv *exec_env, uint32 element_index,
                            uint32 argc, uint32 argv[])
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    // LOG_DEBUG("wasm_runtime_call_indirect from %d\n", gettid());
+#endif
     bool ret = false;
 
     if (!wasm_runtime_exec_env_check(exec_env)) {
@@ -6755,7 +6766,7 @@ wasm_runtime_invoke_c_api_native(WASMModuleInstanceCommon *module_inst,
             wasm_runtime_set_exception(
                 module_inst, "native function throw unknown exception");
         }
-        wasm_trap_delete(trap);
+        // wasm_trap_delete(trap);
         goto fail;
     }
 
@@ -7253,6 +7264,27 @@ wasm_runtime_set_linux_perf(bool flag)
 }
 #endif
 
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+bool
+wasm_runtime_invoke_native_shim(WASMExecEnv *exec_env, void *func_ptr,
+                                const WASMType *func_type,
+                                const char *signature, void *attachment,
+                                uint32 *argv, uint32 argc, uint32 *argv_ret)
+{
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    // Commented because assuming native funcs are not blocking
+    // lightweight_checkpoint(exec_env);
+#endif
+    bool ret =
+        wasm_runtime_invoke_native(exec_env, func_ptr, func_type, signature,
+                                   attachment, argv, argc, argv_ret);
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    // Commented because assuming native funcs are not blocking
+    // lightweight_uncheckpoint(exec_env);
+#endif
+    return ret;
+}
+#endif
 bool
 wasm_runtime_set_module_name(wasm_module_t module, const char *name,
                              char *error_buf, uint32_t error_buf_size)

core/iwasm/common/wasm_runtime_common.h

@@ -1097,6 +1097,14 @@ wasm_runtime_invoke_native_raw(WASMExecEnv *exec_env, void *func_ptr,
                                const char *signature, void *attachment,
                                uint32 *argv, uint32 argc, uint32 *ret);
 
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+#define wasm_runtime_invoke_native wasm_runtime_invoke_native_shim
+bool
+wasm_runtime_invoke_native_shim(WASMExecEnv *exec_env, void *func_ptr,
+                                const WASMType *func_type,
+                                const char *signature, void *attachment,
+                                uint32 *argv, uint32 argc, uint32 *argv_ret);
+#endif
 void
 wasm_runtime_read_v128(const uint8 *bytes, uint64 *ret1, uint64 *ret2);
 
@@ -1202,6 +1210,19 @@ wasm_runtime_detect_native_stack_overflow_size(WASMExecEnv *exec_env,
 WASM_RUNTIME_API_EXTERN bool
 wasm_runtime_is_underlying_binary_freeable(WASMModuleCommon *const module);
 
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+bool
+wasm_runtime_checkpoint(wasm_module_inst_t module_inst, char *file);
+bool
+wasm_runtime_restore(wasm_module_inst_t module_inst, char *file, char *file1);
+#endif
+WASM_RUNTIME_API_EXTERN bool
+wasm_runtime_detect_native_stack_overflow(WASMExecEnv *exec_env);
+
+WASM_RUNTIME_API_EXTERN bool
+wasm_runtime_detect_native_stack_overflow_size(WASMExecEnv *exec_env,
+                                               uint32 requested_size);
+
 #if WASM_ENABLE_LINUX_PERF != 0
 bool
 wasm_runtime_get_linux_perf(void);