Enhance wasm with checkpoint and restore support (bytecodealliance#2333)

- Add wasm_runtime_checkpoint/wasm_runtime_restore API - Support AOT and Classic Interpreter mode checkpoint and debug through OS signal, tested on windows/mac/linux aarch64/x64 - Static instrument the AOT to have the checkpoint and restore switches - Add sub extra library folder for implementing the ckpt-restore - Include extra dependency of yalantinglib Co-authored-by: Aibo Hu <[email protected]> Co-authored-by: kikispace <[email protected]> Co-authored-by: Brian Zhao <[email protected]> Signed-off-by: victoryang00 <[email protected]>
victoryang00 · Jun 7, 2024 · 0f7ec2e · 0f7ec2e
1 parent ac308c4
commit 0f7ec2e
Show file tree

Hide file tree

Showing 69 changed files with 11,516 additions and 304 deletions.
diff --git a/README.md b/README.md
@@ -41,6 +41,7 @@ WebAssembly Micro Runtime (WAMR) is a lightweight standalone WebAssembly (Wasm)
 - [Berkeley/Posix Socket support](./doc/socket_api.md), ref to [document](./doc/socket_api.md) and [sample](./samples/socket-api)
 - [Multi-tier JIT](./product-mini#linux) and [Running mode control](https://bytecodealliance.github.io/wamr.dev/blog/introduction-to-wamr-running-modes/)
 - Language bindings: [Go](./language-bindings/go/README.md), [Python](./language-bindings/python/README.md), [Rust](./language-bindings/rust/README.md)
+- Language bindings: [Go](./language-bindings/go/README.md), [Python](./language-bindings/python/README.md), [Rust](./language-bindings/rust/README.md)
 
 ### Wasm post-MVP features
 - [wasm-c-api](https://github.com/WebAssembly/wasm-c-api), ref to [document](doc/wasm_c_api.md) and [sample](samples/wasm-c-api)

diff --git a/build-scripts/config_common.cmake b/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")

diff --git a/build-scripts/runtime_lib.cmake b/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})
diff --git a/core/config.h b/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

diff --git a/core/iwasm/aot/aot_reloc.h b/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_ */
diff --git a/core/iwasm/aot/aot_runtime.c b/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
@@ -72,6 +75,11 @@ bh_static_assert(offsetof(AOTFrame, sp) == sizeof(uintptr_t) * 5);
 bh_static_assert(offsetof(AOTFrame, frame_ref) == sizeof(uintptr_t) * 6);
 bh_static_assert(offsetof(AOTFrame, lp) == sizeof(uintptr_t) * 7);
 
+bh_static_assert(offsetof(AOTFrame, ip_offset) == sizeof(uintptr_t) * 4);
+bh_static_assert(offsetof(AOTFrame, sp) == sizeof(uintptr_t) * 5);
+bh_static_assert(offsetof(AOTFrame, frame_ref) == sizeof(uintptr_t) * 6);
+bh_static_assert(offsetof(AOTFrame, lp) == sizeof(uintptr_t) * 7);
+
 static void
 set_error_buf(char *error_buf, uint32 error_buf_size, const char *string)
 {
@@ -810,6 +818,7 @@ memory_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
         bh_assert(memory_idx == 0);
         bh_assert(parent->memory_count > memory_idx);
         shared_memory_instance = parent->memories[memory_idx];
+        shared_memory_instance->ref_count++;
         shared_memory_inc_reference(shared_memory_instance);
         return shared_memory_instance;
     }
@@ -1043,6 +1052,15 @@ memories_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
         return true;
     }
 
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    /* Currently we have only one memory instance */
+    bool is_shared_memory = module->memories[0].memory_flags & 0x02 ? true : false;
+    if (is_shared_memory && parent != NULL) {
+        /* Ignore setting memory init data if the memory has been initialized */
+        return true;
+    }
+#endif
+
     for (i = 0; i < module->mem_init_data_count; i++) {
         data_seg = module->mem_init_data_list[i];
 #if WASM_ENABLE_BULK_MEMORY != 0
@@ -1885,6 +1903,7 @@ destroy_c_api_frames(Vector *frames)
 void
 aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst)
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE == 0
     WASMModuleInstanceExtraCommon *common =
         &((AOTModuleInstanceExtra *)module_inst->e)->common;
     if (module_inst->exec_env_singleton) {
@@ -1958,6 +1977,7 @@ aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst)
 #endif
 
     wasm_runtime_free(module_inst);
+#endif
 }
 
 AOTFunctionInstance *
@@ -2242,6 +2262,7 @@ aot_call_function(WASMExecEnv *exec_env, AOTFunctionInstance *function,
         while (exec_env->cur_frame != prev_frame)
             aot_free_frame(exec_env);
 #endif
+        // checkpoint
         if (!ret) {
             if (argv1 != argv1_buf)
                 wasm_runtime_free(argv1);
@@ -2542,6 +2563,9 @@ aot_module_malloc_internal(AOTModuleInstance *module_inst,
     /* TODO: Memory64 size check based on memory idx type */
     bh_assert(size <= UINT32_MAX);
 
+    /* TODO: Memory64 size check based on memory idx type */
+    bh_assert(size <= UINT32_MAX);
+
     if (!memory_inst) {
         aot_set_exception(module_inst, "uninitialized memory");
         return 0;
@@ -2569,7 +2593,7 @@ aot_module_malloc_internal(AOTModuleInstance *module_inst,
 
         if (!malloc_func
             || !execute_malloc_function(module_inst, exec_env, malloc_func,
-                                        retain_func, size, &offset)) {
+                                        retain_func, (uint32)size, &offset)) {
             return 0;
         }
         addr = offset ? (uint8 *)memory_inst->memory_data + offset : NULL;
@@ -2680,7 +2704,8 @@ aot_module_free_internal(AOTModuleInstance *module_inst, WASMExecEnv *exec_env,
                 free_func = aot_lookup_function(module_inst, "__unpin");
 
             if (free_func)
-                execute_free_function(module_inst, exec_env, free_func, ptr);
+                execute_free_function(module_inst, exec_env, free_func,
+                                      (uint32)ptr);
         }
     }
 }
@@ -2857,6 +2882,17 @@ 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;
+        }
+        LOG_DEBUG("func_idx: %d instead of %d of thread %ld\n", rcc->func_index,
+                  func_idx, exec_env->handle);
+        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 +3565,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 +3650,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 +3662,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 +3714,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 +3735,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;
 
@@ -3678,7 +3753,6 @@ aot_free_frame_internal(WASMExecEnv *exec_env)
     if (prev_frame)
         prev_frame->func_perf_prof_info->children_exec_time += time_elapsed;
 #endif
-
     wasm_exec_env_free_wasm_frame(exec_env, cur_frame);
     exec_env->cur_frame = (struct WASMInterpFrame *)prev_frame;
 }

diff --git a/core/iwasm/aot/aot_runtime.h b/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);
 

diff --git a/core/iwasm/common/wasm_exec_env.c b/core/iwasm/common/wasm_exec_env.c
@@ -72,6 +72,7 @@ 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;
+    exec_env->is_checkpoint = false;
 
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT) {
@@ -85,6 +86,10 @@ wasm_exec_env_create_internal(struct WASMModuleInstanceCommon *module_inst,
     wasm_runtime_dump_exec_env_mem_consumption(exec_env);
 #endif
 
+    exec_env->is_checkpoint = false;
+    exec_env->is_restore = false;
+    exec_env->call_chain_size = 0;
+    exec_env->restore_call_chain = NULL;
     return exec_env;
 
 #ifdef OS_ENABLE_HW_BOUND_CHECK

diff --git a/core/iwasm/common/wasm_exec_env.h b/core/iwasm/common/wasm_exec_env.h
@@ -155,6 +155,12 @@ typedef struct WASMExecEnv {
 
     /* The WASM stack size */
     uint32 wasm_stack_size;
+    /* Whether is checkpoint */
+    bool is_checkpoint;
+    /* Whether is restore */
+    bool is_restore;
+    size_t call_chain_size;
+    struct AOTFrame **restore_call_chain;
 
     /* The WASM stack of current thread */
     union {

diff --git a/core/iwasm/common/wasm_memory.c b/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,

diff --git a/core/iwasm/common/wasm_runtime_common.c b/core/iwasm/common/wasm_runtime_common.c
@@ -37,9 +37,11 @@
 #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"
-
 /**
  * For runtime build, BH_MALLOC/BH_FREE should be defined as
  * wasm_runtime_malloc/wasm_runtime_free.
@@ -62,6 +64,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 +3576,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 +5708,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 +6765,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 +7263,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)
-Original file line number
+Diff line change
@@ Expand Up @@
     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);
@@ Expand Down @@